Understanding the Mixture-of-Experts with Nadaraya-Watson Kernel
Authors: 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
Venue: ICLR 2026 Poster (2026)
Keywords: Mixture of Experts
Large Language Models
Foundation Model
Primary area: foundation or frontier models, including LLMs
TL;DR: We rethink the MoE with Nadaraya-Watson Kernel and propose KERN router to replace Softmax router to achieve better performance.

Abstract: 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 Mixture-of-Experts (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 ($\mathrm{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. **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 $\mathrm{KERN}$.

Published as a conference paper at ICLR 2026
UNDERSTANDING THE MIXTURE-OF-EXPERTS WITH NADARAYA- MADARAYA-WATSON KERNEL
Chuanyang Zheng!
Jiankai Sun?
Yihang Gao
Enze Xie
Yuchao Wang
Peihao Wang
Ting Xu
Matthew Chang'
Liliang Ren
Jingyao Li
Jing Xiong"
Kashif Rasul*
Mac Schwager
Anderson Schneider!
Atlas Wang
Yuriy Nevmyvaka'
'Morgan Stanley
Stanford
NUU
*Nvidia HHKU
UT Austin
"CUHK Meta
Microsoft Research
ABSTRACT
Mixture-of-Experts (MoE) has become a cornerstone in recent state-of-the-art large language models (LLMs). Traditionally, MoE relies on Softmax the router score function to aggregate expert output, a designed choice that has persisted from the earliest MoE models to modern LLMs, and 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 princi- pled design choice. In this work, we first revisit the classical Nadaraya- regression and observe that MoE shares the same mathematical formulation as Nadaraya- Watson regression. Furthermore, 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 func- tion, as an alternative to Softmax. We demonstrate that this router generalizes both Sigmoid- and Softmax- routers. Based on empirical observations and established practices in FFN implementation, we recommend the use of ReLU activation and éz-normalization in KERN router function. Comprehensive experiments in MoE and LLM validate the effectiveness of the proposed FFN-style router function KERN.
show that both feed-forward neural
INTRODUCTION
Recent years have witnessed remarkable progress in Large Language Models (LLMs) (Brown et al., 2020; Ouyang al., 2022; Touvron al., 2023), driven primarily by the exponential growth of training data and model parameters. With the mixture of experts (MoE), there is great progress in language modeling (Fedus al., 2022; Puigcerver al., 2024; Jiang al., 2024; Meta; Liu et al., 2024a; Team et al., 2025) and computer vision (Riquelme et al., 2021; [in et al., 2024b). The MoE architecture (Jacobs et al., 1991; Shazeer al., 2017: Roller et al., 2021) has emerged as an efficient alternative that allows parameter scaling while maintaining manageable computational requirements. The successful integration of MoE with Transformer architectures (Vaswani et al., 2017) has led to the development of exceptionally large yet efficient language models (Dai et al., 2024; Jiang et al., 2024: Shen et al., 2024; Wei et al., 2024; Liu et al., 2024a), demonstrating the tremendous potential of this approach.
critical and widely adopted design choice in modern MoF. architectures is the use of the Softmax function as the core routing mechanism. This approach, prominently featured in large-scale mod- els (Lepikhin et al., 2021; Jiang et al., 2024; Liu et al., 2024a; Team et al., 2025), has effectively become the de facto standard for state-of-the-art systems. The function Softmax naturally induces a probability distribution on the available experts. This property ensures that the routing weights for each token sum to one, promoting balanced and interpretable allocation. However, despite its prevalence and intuitive appeal, the theoretical justification for its exclusive dominance remains
"Contact Email: @yzhengme@gmaile gmail.com

---

Published as a conference paper at ICLR 2026
somewhat unclear. Recently, Sigmoid has been proven to be a better router function (Nguyen et al., 2024a), which is also investigated and adopted as an alternative router score function by DeepSeek (Dai et al., 2024; Liu et al., 2024a). Their findings suggest that Sigmoid-basee routing function performs effectively in MoF.
In this work, we revisit the fundamental design principles of MoE routing by re-examining it through the statistical lens of the Nadaraya-Wart Watson regression estimator (Nadaraya, 1964; Watson, 1964). We propose a novel interpretation: the router's output for given input token can be viewed set of dynamic kernel weights assigned to each expert. Each expert, in turn, acts as value function, producing an aggregated output. This perspective is further reinforced by the architectural parallels within the Transformer ( (Vaswani et al., 2017). We posit that the router's s computation is analogous to the first linear layer of a standard forward network (FFN), which projects the input into higher- dimensional space and can be interpreted as calculating a set of unnormalized scores or weights. The experts subsequently play the role of the second FFN layer, which operates on these weighted features to produce the final values (Geva et al., 2020). Inspired by structural similarities between Watson regression, and FFN, we introduce a new class of simple yet effective router functions for MoE. Our primary proposed method defines an FFN-style router function, which generalizes both Softmax- and Sigmoid- widely recognized FFN setups, adopt the ReLU activation and computationally lightweight f2-normalization in the router function. This modification ensures that the magnitude of the MoF output is invariant with the number of experts, leading more balanced expert participation and improved training stability without enforeing probabilistic simplex constraint. Our key contributions are summarized as follows:
MoE, Nadaraya- Wadaraay--aatoon
router functions. To align well with the practical and
A Novel Perspective:
We reframe the MoE layer through the lens of the Nadaraya- Nadaraya-Watson Watson regression, interpreting it as a generalized FFN, providing a more flexible and principled view of expert aggregation.
KERN Router Function: Motivated by the perspective of structure similarity, we pro- pose Kernel Inspired Router with Normalization (KERN), a new family of simple yet effective FFN-style router functions. By introducing widely adopted ReLU activation and fz-normalization, KERN promotes balanced expert utilization and stable training without the constraints computational profile of Softmax, and crucially, without introducing any additional parameters or significant overhead.
Extensive Empirical Validation:
We conduct a comprehensive evaluation of KERN across wide range of experimental setups, including varying model scales, sequence lengths, training dataset sizes and domains, and sparsity coefficients.
2 RELATED WORK
Large Language Models. With the inspiration of the language model scaling law (Kaplan et al., 2020), LLMS (Touvron et al., 2023; Achiam et al., 2023; Jiang et al., 2024; Liu et al., 2024a; Yang et al., 2025; Team et al., 2025; Comanici et al., 2025) have shown remarkable capabilities in a wide range of open-ended tasks, marking significant progress toward achieving general artificial intelligence. With the Transformer architecture (Vaswani et al., 2017), LLMs achieve significant performance in various areas, including reasoning (Achiam et al., 2023; Liu et al., 2024a; Team et al., 2025), language- visuall model (Liu et al., 2023a; Jin et al., 2024a; Riquelme et al., 2021), language-audio (Yang et al., 2023; Rouditchenko et al., 2025)
and so on.
Hixturee-f-Expperts The MoE (Jacobs et al., 1991) is proposed reduce the active parameters and aggregate the outputs of several models to reduce the training cost and empower expressiveness. With the development of L.LMs, the MoE becomes increasingly attractive and dominant applications of large-scale tasks (Achiam et al., 2023; Meta), where they must balance the load of experts (Lewis et al., 2021; Roller et al., 2021; Dai et al., 2024). The MoE originally presents its ability in the machine translation tasks (Shazeer et al., 2017). Later, Gshard (Lepikhin et al., 2021) proposes a more efficient implementation on parallel devices. To further improve the efficiency, Switch Transformer (Fedus et al., 2022) alternatively uses single expert for one token prediction. Recently, a zero-cost expert (Jin et al., 2024b) is introduced, where the expert does not involve computation via skip connections. We notice that most of these works utilize Softmax the router function.
2

---

Published as a conference paper at ICLR 2026
Feed-Forward Network. FFN represents the standard neural network architecture, with origins tracing back to the early development deep learning. Numerous studies have examined different activation functions and normalization techniques to enhance their expressiveness and training stability (Householder, 1941; Fukushima, 1980; 2007; Hendrycks & Gimpel, 2016). More recently, research has shifted toward understanding the role of FFNs within Transformer models, where they are often interpreted as form of static memory, in contrast to the dynamic memory provided by attention mechanisms. From this perspective, Transformers can be viewed as integrating both static and dynamic memory, each contributing distinct modes of information processing (Liu et al., 2023b; Zhong et al., 2025).
METHOD
In this section, we first introduce the well-knoww Nadaraya- -attaasyy-attonnnnno Watson regression and then compare its mathematical formulation with that of MoE. Motivated by their structural similarity, we reinterpret the MoE as large FNN. Inspired by the new perspective of interpretation and the well-recognized FFN setups, we design an FFN-style router function of MoE, namely, KERN, which is equipped with ReLU activation and ez-normalization, We also analyze the relationships and advantages of KERN, compared to widely recognized Softmax and Sigmoid router functions.
3.1 NADARAYA- WATSON REGRESSION
The Nadaraya- Watson estimator predicts the output for an input x: by assigning weights to training samples {(x,, yi) W according to their similarity to
fNw (z) = = K(x,w:;) y:
N
(1)
j== 2()
where ((::)) () a kernel function measuring the similarity between two points. The most widely used choice is the Gaussian kernel formulated as K (u, v) exp(-Ilu u /20²), where the bandwidth a (standard deviation of the Gaussian distribution) controls the smoothness of the estimator.
In practice, the bandwidth is typically unknown and treated as hyperparameter. A more flexible approach is to regard a as a trainable parameter to optimized. Equivalently, we define a param- eterized kernel as K uu, 2; w) exp(-wllu w/(2/2), where 0 is learnable weight. The corresponding estimator becomes
K(x, a;; f2w( = Σ E K(
N
yi-
This parametric formulation allows the model to adapt the kernel bandwidth during training, improv- ing flexibility and performance in practice. Moreover, the idea naturally extends beyond Gaussian kernels that one can generalize to learnable kernel of the form (u, 2; w), parameterized by a vector w.
3.2 FFN AS PARAMETRIC NADARAYA- WATSON REGRESSION
The output layer of a standard FFN admits two potential formulations
FFN(#) = Σ (LN ((w,, (x))))
V
(3)
Adaptive kemel weight
Value
FFN(#) Σ LN (o ((w,, (())))
v:
Adaptive kemeel weight
Value
where is the activation function in the FFN, ((x)) denotes the hidden representation input to the output layer, and V = "0,,...,y,, are the output-layer weights. Here, LN(-) denotes layer normalization.
3

---

The structure in Equation (6) closely resembles Nadaraya- Mattaray--Watonn Watson regression in Equation (2) that the router weight 9m. (x) can be viewed kernel function K [xx, w.),, while each expert output (x) corresponds to an observation ym being aggregated. Recall from the previous section that we Published as a conference paper at ICLR 2026
Comparing Equation (4) with the adaptive Nadaraya- Watson estimator in Equation (2), we that the FFN implicitly defines a parameterized FFN-style kernel function
(xx (w, b,)) = ø ((w,, ((x))),
(5)
where the normalization is applied after the kernel function, the role of the labels yi played by the value vectors Vi, and Φ is transformation function. In this analogy, the normalization step in Equation (2) corresponds to --normalization IN (x: || This observation motivates a natural
generalization of replacing the normalization LN (-) in Nadaraya- Watson regression with a more commonly used (2-normalization in FFN. This perspective provides mathematical interpretation of the FFN as a special instantiation of parametric Nadaraya- Madaraya-Watson Watson regression.
3.3
MIXTURE-OF-ENTE EXPERTS EXPPERTS AS PARAMETRIC NADARAYA WATSON REGRESSION
The MoE model combines multiple expert networks {Em { m
through a router g(:):
MoE(x) Σ Im (2:
(6)
where the router g(x) (z)( gM (a)) admits Softmax outputs:
exp( exp((Wum: x²)
9m (x)
ΣØ= exp((ux; x))
M
generalized the kernel function to the FFN-style form given in Equation (5). Under this perspective, the MoE in Equation (6) can be interpreted and designed as large network that aggregates expert networks {Em (x)) m= via such an FFN-style kernel function given by:
Im (a) Φ (LN ((wm: ((x)))))
3.4 KERN ROUTER FUNCTION
Let (x)) Rd denote the representation that feeds the router. We introduce novel router function defined in Equation (7), namely, the kemellinspired router with normalization (KERN), that instan- tiates the FFN-style router with linear projection followed by (i) & fg-normalization, (ii) ReLU
activation, and (iii) an optimal learnable global scaler:
s <xx) W. (xx()) bs:
s(x))
8 s(x) |s(x)]1 +
(8)
r(x) RcLUUs(x:)).
((x) 7 -x)).
where W, Mxd and b, E R are the router parameters, a small constant that guards against division by zero, and is a learnable scalar initialized to 1. The normalization step keeps the scale of the logits invariant to the number of experts M, while the ReLU activation preserves sparsity without resorting to exponential functions. We further discussed the effect of ReLU in Appendix J. During inference and training, we retain only the top-k routed experts:
(xx TopKIndices(g(x)))
(9)
9m (x) = 9m (x) 1mm E Tx(x)),
(10)
M
MOEKERN (x) Σ
9m 9m(x) (m)).
(11)
m=1
4

---

Published as a conference paper at ICLR 2026
Because KERN does not project the router outputs onto the probability simplex, do not perform an additional & rescaling; the magnitude of 9m (x) is instead controlled by the global scale y and the l2 constraint. This simple construction matches the inductive biases of standard FFNs while avoiding the gradient saturation issues of exponential routers.
Comparisons to existing router functions. From this viewpoint, the standard MoE with Softmax router corresponds to an FtN-style router where & fn-normalization is applied through LN(-) and the exponential function serves as the activation. Interestingly, recent work has shown that replacing the Softmax with a Sigmoid router yields improved performance (Nguyen et al., 2024a; Chi et al., 2022; Wang et al., 2024b;a). It also admits an FFN interpretation that the router function reduces to a Sigmoid activation o(-) without layer normalization. Hence, our interpretation frames MoE routing as general framework that encompasses the most widely adopted router functions. However, under this formulation, the commonly used Softmax and Sigmoid routers appear atypical when compared with standard FFN activations and layer normalization deep learning. This observation motivates us to explore and design the router function KERN, which is more consistent with the classical FFN paradigm, as discussed in the next paragraph.
Practical setups for the FFN-style router function KERN. To better align with common practices in deep learning, we propose adopting router functions that mirror typical FFN designs, namely using ReLU activation o(.) combined with widely-adopted ly-normalization LN(.). This choice is motivated by the following observations. First, exponential-type activations are rarely used in modern architectures, they tend to be highly sensitive input values, leading to rapid value explosion and gradient vanishing, In contrast, ReLU-type activations, or even linear outputs without nonlinearity, are far more common in practice, providing numerical stability and robustness during training. Second, although all vector norms are mathematically equivalent in finite-dimensional spaces, l2- 22-normalization remains the dominant choice in deep learning. Specifically, & n-omaalizationn stabilizes the variance of vectors, ensuring scale consistency and stable computation regardless of model size.
Advantages of the KERN router function. First, the gradient vanishing problem inherent in Softmax and Sigmoid (exponential-tyy) router functions can be alleviated by adopting the proposed FFN-style router (KERN) with appropriate activation functions. Prior studies have highlighted that Softmax and Sigmoid activations often suffer from saturation, that small values push experts toward near-inactivity resulting in negligible gradient updates. Intuitively, if an expert stays at almost-zero routing weight, the vanishing gradient problem can trap it in this poor state, preventing meaningful updates improvements. In contrast, the proposed KERN reduces this risk. The gradients vanish less severely, ensuring that even less active experts still receive updates, therefore, promoting better expert utilization and training dynamics. Second, nommaaaaliotio in KERN preserves the variance of the MoE output at a constant scale, Since experts are independently and properly initialized, we may assume that the outputs of the M experts, {m (x)) M m=1 are independent and have bounded norm (i.e., |m) O())). Under this assumption, the MoE with KERN at initialization satisfies
M
E (MOEKEnN(2)1))
=
Σ
Imn (ax) Fm (xx)
M
(Im (x))² E Em (2)
m=1
O() (=)(=)===011
The final equality holds for most commonly used activation functions & (e.g., ReLU, Lea,yReLU, Tanh, GeLU) when applied within the proposed FFN-style router function, where KERN adopts ReLU activation. This result demonstrates that ez-oormalization maintains the MoE output at constant scale, thereby ensuring stable network computations and training. Such stability is consistent with the initialization principles commonly adopted in deep neural networks, e.g., Kaiming initialization.
5

---

Published as a conference paper at ICLR 2026
4 EXPERIMENT
Baseline. We compare the proposed KERN with the Dense model and MoE with other router functions. To be specific, we evaluate KERN against range of routers, including Softmax, Sigmoid, and Tanh. For the MoE and Dense model, they have the same active parameters. Additionally, for the MoE model, the ratio of active parameters and total parameters is 8, where there are 64 experts in total and active experts. We utilize more than 8 active experts, as recent works propose using a larger number of active experts (Liu et al., 2024a; Team et al., 2025).
Datasets. Our analysis involves training language models on the Arxiv and Books3 datasets, which are frequently used benchmarks for evaluating model performance (Press et al., 2022). Moreover, we train the model on the large-scale dataset Fin Web-Edu (Penedo et al., 2024; Lozhkov et al., 2024) and evaluate on downstream datasets, including ARC (Clark et al., 2018), HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020), ScIQ (Welbl et al., 2017), and WinoGrande (Sakaguchi et al., 2021)
Experiment settings. Initially, compare KERN with other baselines at training lengths 512 and 1024, using decoder-only Transformers (Brown et al., 2020) with model size 125M, whose configuration is shown in Appendix D. Subsequently, sizes, specifically 350M and 2.7 B. Finally, we analyze routers and MoE models by various active experts while holding the active ratio fixed, and we examine the effect of sparsity.
evaluate the performance of larger model
4.1
COMPARISONS WITH BASELINES
Arxiv Dataset (Length 512)
Books3 Dataset (Length 512)
Dana Sipred
Dama
4.2
Sprvil Tarb
Validacion Less
2.2
Validation
3.8
2.0
3.6
3.4
1.8
10K
15K
20K
25K 30K 35K 40K
45K 50K
5k
10K
15K
20K
25K 30K 35K 40K 45K 50K
Validation Step
Validation Step
Arxiv Dataset (Length 1024)
Books3 Dataset (Length 1024)
2.2
Devse Sofmra Twlk
Dowc
Validation Lass
Validation Loss
3.8
KIN
KERH
2.0-
3.6
3.4
3.2-
5k 10K
15K 20K 25K 30K 35K 40K 45K 50K Validation Step
5k
10K
15K 20K 25K 30K 35K 40K 45K 50K Validation Step
Figure The performance of different methods on the Arxiv and Books3 dataset, with model parameter 520M, activated parameter 125M, training lengths of 512 and 1024.
KERN achieves superior performance across different datasets. We validate our method on the Arxiv and Books3 datasets in Figure 1. On Arxiv with training length 512, the Dense model (125M) reaches losses of 2.3396 and 1. 1.0925 at steps 5,000 and 50,000, respectively. The Softmax router shows a higher initial loss (2.4586) but a better final loss (1.8781), suggesting potentially slower convergence rate but strong final performance. Our KERN method achieves the best results at both checkpoints (2.3975 and 1.8291). A similar trend is observed on the Books3 dataset. The Dense model records 4. 1460 (step 5,000) and 3.4429 (step 50,000). The Softmax router achieves 4.301 I and 3.3882. Once again, KERN delivers the best performance, with losses of 4.2165 and 3.3080 at steps 5,000 and 50,000, respectively. Therefore, regardless of the training dataset, KERN consistently achieves state-of-the-are performance.

---

Published as a conference paper at ICLR 2026
KERN achieves superior performance across varying training lengths. We further evaluate model performances among various router functions using a context length of 512 on the Books3 dataset. The baseline Dense model achieves loss of 3.4429. Among these MoE routers, Sigmoid (3.3206), Tanh (3.3388), and Softmax (3.3882) all outperform the Dense baseline. The proposed KERN method achieves the best performance with loss of 3.3080. When the context length is increased to 1024, the performance ranking remains consistent: the Dense model attains loss of 3.2454, while Softmax, Tanh, and Sigmoid achieve 3.1714, 3.1224, and 3.1031, respectively. Notably, KERN again achieves the lowest loss 3.0914, demonstrating its robustness across different training lengths.
KERN achieves superior performance with longer context lengths. The advan-
Books3 Dataset (Length 2048)
tage of KERN is further demonstrated at a longer context length of 2048 in Figure 2. The baseline Dense model achieves a loss of 3.1249. The Softmax router shows an improvement with loss of 3.0442, while Sigmoid (2.9635) and Tanh (2.9868) perform better still. The proposed KERN method achieves the best perfor- mance overall, with lowest loss of 2.9535. These results confirm that KERN main- tains its effectiveness and superiority as the training length increases.
3.8
Validation Lod
3.6
3.4
3.2
3.0-
5k
10K
15K
20K 25K
35K 40K 45K 50K
Validation Step
Figure 2: The performance on training length 2048.
4.2
THE PERFORMANCE ON LARGE LANGUAGE MODELS
1.7B Parameter (Activated 350M, Length 512)
6.9B Parameter (Activated 1.38, Length 512)
Dresa
4.2
Sofmee Tam KoW
4.0
Val dation Loss
Tae
3.8
3.6
3.6
3.4
3.4
3.2
3.2
5k
10K
15K
20K 25K 30K 35K 40K 45K 50K
5K
10K
15K
20K 25K 30K 35K 10K 45K 50K
Validation Step
Validation Step
Figure 3: The performance of different methods on the Books3 dataset, with active model size 350M and 1.3B.
The performance gap provided by KERN is maintained for large models. For Figure 8, while increasing the active model size from 125M to 1.3B parameters reduces the loss for all baseline methods, KERN consistently achieves the best performance every scale. At 350M parameters, its loss of 3.2188 88 is lower than that of the comparable Softmax (3.2709) and Dense (3.3500) models. This lead is extended at the 1.3B scale, where KERN's loss 3.1241 significantly outperforms the Softmax (3.1814 18)4) and Dense (3.2219) results. This demonstrates that KERN is not only effective but is particularly advantageous for training larger-scale models.
4.3
THE EFFECT OF GRANULARITY
KERN consistently outperforms Softmax across expert counts. Figure 4 provides compelling evidence for the practical superiority of KERN. When evaluated against the standard Softmax baseline, our method achieves higher performance regardless of the number of experts activated during inference-a parameter we varied from 4 to 32. This demonstrates the robustness of our method irrespective of the specific capacity used during inference.
KERN achieves better performance than other routers with small granularity (e.g., more experts, smaller expert size). As shown in Figure 7, with 256 experts, 8 active experts, and an expert intermediate size is 96, the Softmax achieves 4.3139 loss at evaluation step 5K and 3.4150 loss
7

---

Published as a conference paper at ICLR 2026
Different Granularity (Arxiv)
Different Granularity (Books3)
Setmman 32 4.359
Setman 4.20
2.4
Setiman 259-2-0
Validation Loss
Settman
4.0
RAke 22
Validation
KSSI 3.0 28-353
3.6
1.8
5k 10K 15K 20K 25K 30K 35K 40K 45K Validation Step
50K
5K 10K
15K 20K 25K 30K 35K 40K 45K 50K Validation Step
Figure 4: The performance of different methods on the Arxiv and Books3 dataset, with different active expert numbers and the same active parameter number. The Softmax 64 - - 384 suggests that the router function is Softmax, and there are 64 experts, active experts, and each expert's intermediate size is 384.
evaluation step 50K. The Sigmoid and Tanh achieve 4.2924 loss and 4.2435 loss at the evaluation step 5K, and Sigmoid and Tanh achieve 3.3302 loss and 3.3276 loss at the evaluation step 50K. The KERN achieves the best performance 4.2294 loss at evaluation step 5K and 3.2962 loss evaluation step 50K. Therefore, KERN achieves better performance than other routers with small granularity.
4.4 THE EFFECT OF SPARSITY
Different Sparsity Number (Arxiv)
Different Sparsity Number (Books3)
Cotma-332--04
Soltmio-238
2.4
4.2
5obmmn-1236-30
Saftmeo 8284
Validation Less
50tmax25556-3
Validation Less
Sottman--36 26
2.2
AER9-04-830
1EIN-66-334
KLR--126-330
3.8
2.0
3.6
3.4
1.8
5k 10K
15K 20K 25K 30K 35K 10K 45K 50K Validation Step
5K
10K
15K 20K 25K 30K 35K 10K 45K 50K Validation Step
Figure 5: The performance of different methods on the Arxiv and Books3 dataset, with different total expert numbers and the same active parameter number. The Softmax - 64 - 8 384 suggests that the router function is Sofumax, and there are 64 experts, 8 active experts, and each expert's intermediate size is 384.
KERN outperforms Softmax across all sparsity levels. As evidenced by Figure 5, KERN
achieves a lower loss than the Softmax baseline for every total number of experts tested (32 to 256) on both the Books3 and Arxiv datasets. On Books3, KERN's loss (3.3487, 3.3080, 3.2820, 3.2672) consistently superior Softmax's (3.3981, 3.3882, 3.3817, 3.3761). 3.3661, This trend holds on the Arxiv dataset, where KERN's results (1.8466, 1.8291, 1.8195, 1.8141) are consistently better than those of Softmax (1.8835, 1.8781, 1.8738, 1.8754). This demonstrates that the performance gain of KERN is robust to changes in model sparsity.
KERN achieves better performance than other routers with large sparsity. As shown in Figure 6, with 256 experts, 8 active experts, and an expert intermediate size of 384, the Softmax achieves 4.3059 loss at evaluation step 5K and 3.3761 loss at evaluation step 50K. The Sigmoid and Tanh achieve 4.2924 loss and 4.2435 loss
Large Sparsity Performance (Books3)
Validztion Loss
4,0
3.8
3.6
at the evaluation step 5K, and Sigmoid and
10K
15K 26K 25K 30K 35K 40K 45k 50k Validation Stap
Tanh achieve 3.2760 loss and 3.2972 loss at the evaluation step 50K. The KERN achieves the best performance 4, 1926 loss
Figure 6: The performance with 256 experts, 8 active experts and each expert's intermediate dimension is 384.
8

---

Published as a conference paper at ICLR 2026
at evaluation step 5K and 3.2672 loss at evaluation step 50K. Therefore, KERN achieves better performance than other routers with large sparsity.
4.5 THE PERFORMANCE ON LARGE-SCAA PRETRAIN DATASET
Table I: Main language modeling results against different methods. All models are trained on the same subset of the FineWWeeEdl dataset (Penedo al., 2024; Lozhkov et al., 2024) with the GPT-2 tokenizer.
Model
ARC-E
ARC-C
Hellaswag
PIQA
ScIQ
Winograde
Avg
520M (Active 125M) params
Dense
47.35
20.48
31.59
66.21
73.60
51.85
48.51
Softmax
49.49
20.39
34.81
69.42
75.50
49.64
49.88
Tanh
51.60
21.59
35.90
69.91
76.90
53.28
51.53
Sigmoid
51.89
20.99
37.11
70.78
78.50
51.54
51.80
KERN
53.32
21.67
37.12
70.89
77.80
52.01
52.14
1.7B (Active 350M) params
Dense
51.26
22.10
35.07
70.13
77.60
50.12
51.05
Softmax
51.94
22.78
37.65
70.13
79.70
52.57
52.46
Tanh
55.51
23.21
40.30
72.03
81.60
52.41
54.18
Sigmoid
56.61
23.81
40.78
72.74
80.40
53,99
54.72
KERN
56.48
24.40
40.68
73.61
82,00
53.59
55.13
6. 6.9B (Active 1.3B) params
Dense
58.59
24.15
42.36
72.85
82.90
56.11
Softmax
42.29
73.18
84.70
55.72
56.49
Tanh
45.01
73.29
85.20
58.04
Sigmoid
62.33
27.47
45.52
74.43
84.70
58.55
KERN
45.95
75.19
84.90
58.88
Downstream Evaluation. We evaluate performance standard benchmarks, including ARC (Clark al., 2018), HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020), ScIQ (Welbl al., 2017), and WinoGrade (Sakaguchi et al., 2021), using the 1m-eva bation-ha rness (Gao et al., 2024) codebase, The evaluation metric is the accuracy. We train the model with 50K steps with training length 1024 and training tokens 50B. The model sizes are 520M (active 125M), 1.7B (active 350M), and 6.9B (active 1.3B). We display the zero-shot evaluation results models here in Tables 1.
With the same active parameter, the KERN is always better than the routers, from small model size (e.g., 125M active) to larger model size (e.g., 1.3B active). At 520M model size (125M active), KERN achieves an average performance of 52. 14, surpassing Dense (48.51), Softmax (49.88), Tanh (51.53), and Sigmoid (51.80). With 1.7B model size (350M active), it scores 55.13, outperforming Dense (51.05), Softmax (52.46), Tanh (54.18), and Sigmoid (54.72). Similarly, at 6.9B (1.3B active), it reaches 58.88, exceeding Dense (56.11), Softmax (56.49), Tanh (58.04), and Sigmoid (58.55). Therefore, KERN is always better than the routers, from small model size to larger model size.
With the same active parameters, the performance gap between KERN and Softmax is com- parable to that between Softmax and Dense model. For a 520M model (125M active), the performance gap between KERN and Softmax is 2.26. compared to the 1.37 gap between Softmax and Dense. With a 1.7B model (350M active), the KERN-Soft max gap widens to 2.67, while the Softmax- Doftmax-Dense gap 0.61. At the 6.9B scale (1.3B active), the KFRNNSSfttax gap remains significant at 2.39, vastly exceeding the 0.38 gap between Softmax and Dense. The MoE model achieves significant performance gain over the Dense model. Since KERN achieves an even larger gain at zero additional cost, it should be critical component for model sparsity.
5 CONCLUSION
In general, the use of the Softmax function has been the de facto standard for generating router scores in MoE models. In this work, we challenge this convention by recasting MoE routing through the novel lens of the Nadaraya- Watson estimator. Motivated by this perspective, introduce the

---

Published as a conference paper at ICLR 2026
novel KERN router function for MoE, an FFN-style kernel function with ReLU activation and lz 2g-normaiizaton. We extensively validate the efficacy of these functions through comprehensive experiments acToss varying model scales, sequence lengths, and, most significantly. in large-scale pre-training followed by downstream task evaluation. Our empirical results demonstrate that these simpler alternatives are not only viable but often match exceed the performance of Softmax-bass routing. We believe this work opens new direction for router design and anticipate that KERN will serve as strong baseline and a potential substitute for Softmax in future MoE architectures.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.065: 303008774, 2023.
Stella Biderman, Hailey Schoelkopf. Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference Machine Learning, pp. 2397-2430. PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 13:1877-1901, 2020.
Zewen Chi, Li Dong, Shaohan Huang, Damai Dai, Shuming Ma, Barun Patra, Saksham Singhal, Payal Bajaj, Xia Song, Xian-Ling Mao, et al. On the representation collapse of sparse mixture of experts. Advances in Neural Information Processing Systems, 35:34600-34613, 2022.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung. Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240) 2023.
Aidan Clark, Diego de Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference machine learning, pp. 4057-4086. PMLR, 2022.
Christopher Clark, Kenton Lee, Ming-Wor Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lan- guage Technologies, Volume 1 (Long and Short Papers), pp. 2924-2936, Minneapolis, Min- nesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v 1/N13///1/1111100000 URL org/ N19 1300/.
ht
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv: 0503. 2018
Gheorghe Comanici, Eric Bieber, Mike Schaekermann, Ice Pasupat, Noveen Sachdeva, Inderjit Dhillon, Marcel Blistein, Ori Ram, Dan Zhang, Evan Rosen, et al. Gemini 2.5: Pushing the frontier with advanced reasoning, multimodality, long context, and next generation agentic capabilities. arXiv preprint arXiv:2507.06261, 2025.
Damai Dai, Chengqi Deng, Chenggang Zhao, R.x. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie. Y.k. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. DeepSeekMoE: Towards ultimate expert specialization mfxture-of experesss language models. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings the
10

---

Published as a conference paper at ICLR 2026
62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), PP. 1280-1297, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: org/ .al-2on.clloog
10.11663/w1120222.allllozzg 70/-
URL htttps://aclattthtttttt
Jacob Devlin, Ming-Wai Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association, for compulational linguistics: human language technologies, volume (long and short papers), pp. 4171-4186, 2019.
Nan Du, Yanping Huang, Andrew M Dai, Simon Tong, Dmitry Lepikhin, Yuanzhong Xu, Maxim Krikun, Yanqi Zhou, Adams Wei Yu, Orhan Firat, et al. Glam: Efficient scaling of language models with mixture-of-experts. In International conference on machine learning, pp. 5547-5569. PMLR, 2022.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research, 23(120):1-33) 2022.
Kunihiko Fukushima. Ncocognitron: A self-organizing neural network model for mechanism of pattern recognition unaffected by shift in position. Biological cybernetics, 36(4): 202, 1980.
Kunihiko Fukushima. Visual feature extraction by a multilayered network of analog threshold elements. IEEE Transactions on Systems Science and Cybernetics, 5(4):322-333, 2007
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.0002 2020.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. URL https:tt/......... /ttps://zenneeee wenodo.orggg records / /00/22208666 2268660..
Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forward layers are key-value memories. arXiv preprint 101V:22012.149133 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08411 08415, 2016.
Alston S Householder. A theory of steady-state activity in nerve-fiber networks: I. definitions and preliminary lemmas. The bulletin of mathematical biophysics, 3(2):63-69, 1941.
Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey Hinton. Adaptive mixtures of local experts. Neural computation, 3(1):79-87, 1991.
Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts. arXiv preprint 2:XX::240100-0088. 2024.
Peng Jin, Jinfa Huang, Pengfei Xiong, Shangxuan Tian, Chang Liu, Xiangyang Ji, Li Yuan, and Jie Chen. Video-text as game players: Hierarchical banzhaf interaction for cross cross-modal representation learning. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 2472-2482, 2023.
Peng Jin, Ryuichi Takanobu, Wancai Zhang, Xiaochun Cao, and Li Yuan. Chat-univi: Unified visual representation empowers large language models with image and video understanding. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 13700-13710, 2024a.
11

---

Published as a conference paper at ICLR 2026
Peng Jin, Bo Zhu, Li Yuan, and Shuicheng Yan. Moe++: Accelerating mixture-of-exppee methods with zero- computatmtttioon experts. arXiv preprint arXiv:2411.07348, 2024b.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0883 201::2001.0866), 2020.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-666, 2019.
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. RACE: Large-scale ReAding comprehension dataset from examinations. In Martha Palmer, Rebecca Hwa, and Sebastian Riedel (eds.). Proceedings of the 2017 Conference on Empirical Methods Natural Language Processing. pp. 785-794, Copenhagen, Denmark, September 2017. Association for Computational Linguistics. doi: 10. 10.18653/11/D17-10022 URL hwtps::tt//....... / antho logy www.gy.org/iii1 2tgg/D1--182//
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2021. URL https: / www.reveeeeeevee net/ / forum? dd-irreeddwwwwee 7XHTmYb.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. Base layers: Simplifying training of large, sparse models. In International Conference Machine Learning, pp. 6265-6274. PMLR, 2021.
Bin Lin, Zhenyu Tang, Yang Ye, Jiaxi Cui, Bin Zhu, Peng Jin, Jinfa Huang, Junwu Zhang, Yatian Pang, Munan Ning, et al. Moe-Ilava: Mixture of experts for large vision-language models. arXiv preprint 21X::24011119477 2024a.
Bin Lin, Yang Ye, Bin Zhu, Jiaxi Cui, Munan Ning, Peng Jin, and Li Yuan. Video-LLaVA: Learning united visual representation by alignment before projection. Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.), Proceedings of the 2024 Conference Empirical Methods in Natural Language Processing, pp. 5971-5984, Miami, Florida, USA, November URL
2024b. Association for Computational Linguistics. doi: https:: / aacl antho logy org/ 2024 emn 1 p-maa i in. 342/.
Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024a.
Dongyang Liu, Renrui Zhang, Longtian Qiu, Siyuan Huang, Weifeng Lin, Shitian Zhao, Shijic Geng, Ziyi Lin, Peng Jin, Kaipeng Zhang, Wenqi Shao, Chao Xu, Conghui He, Junjun Hc, Hao Shao, Pan Lu, Yu Qiao, Hongsheng Li, and Peng Gao. SPHINX-x: Scaling data and parameters for a family of multi-modal large language models. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (eds.), Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 12400-32420 PMLR, 21-27 Jul 2024b. URL https: /t/pppeeeenngsss
www.rrpresss/eeeee /itess/vie v2 3 1055//iui Liuu40o co .html.
Haotian Liu, Chunyuan Li, Qingyang Wu, and Yong Jae Lee. Visual instruction tuning. Advances in neural information processing systems, 36:34892-3491 2023a.
Haotian Liu, Chunyuan Li, Yuheng Li, and Yong Jae Lee. Improved baselines with visual instruction tuning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 26296-26306, 2024c.
Zeyu Liu, Tim Dettmers, Xi Lin, Veselin Stoyanov, and Xian Li. Towards a unified view of sparse feed-forward network in pretraining large language model. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 15038-15061, Singapore, December 2023b. Associa- tion for Computational Linguistics. doi: 10. 18653/v www.:mmmmmlmmm..9999 URL https: //aclantho ogy. org/ 2023. emn emnlp-mat in. 930/.
12

---

Published as a conference paper at ICLR 2026
Anton Lozhkov, Loubna Ben Allal, Leandro Werra, and Thomas Wolf. Fineweb-edu: the finest
collection of oducational content, 2024. URL https::/hhbt httppss:///ggggggca....oo cot stt.://aacecccattatttsss/sss/s//
datasets/
AI Meta. The llama herd: The beginning of new era natively multimodal ai innovation. https:/ / /ai. meta. meeta.combblog com/blog/ 1lama-4: 4-mult wwweimoaaaanneeiiieeeeeeeeeeeeee 4-7-2025.
Accessed:
Elizbar A Nadaraya. On estimating regression. Theory of Probability & Its Applications, 9(1): 141-142, 1964.
Huy Nguyen, Nhat Ho, and Alessandro Rinaldo. On least square estimation in softmax gating mixture of experts. arXiv preprint arXiv:2402.02952, 2024a.
Huy Nguyen, Nhat Ho, and Alessandro Rinaldo. Sigmoid gating is more sample efficient than softmax gating in mixture of experts. Advances in Neural Information Processing Systems, 37: 118357-118388 2024b.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances neural information processing systems, 35:27730- 27744, 2022.
Guilherme Penedo, Hynek Kydlfček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL https:/ /wtp//ppnneeew.eeoe//e net/ www.//eeet///./...e/ 2dG.
Ofir Press, Noah Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. In International Conference Learning Representations, 2022. URL
btttps://ppenreveee wntpeeiwwwwweeeeeeeeeeeee.eeeeee http::/weweeeeeeeeeeeeeee Lormm?1--R8s@@pGGoV
Joan Puigcerver, Carlos Riquelme Ruiz, Basil Mustafa, and Neil Houlsby. From sparse to soft
mixtures of experts. In The Twelfth International Conference on Learning Representations, 2024.
URL hwtph:tpp////. http/ppeeeeeeeeeeeeeeeeeeeeeeeee eeeerveewww net forum? ? id= jxpsA joapssj71ttt j7ltE.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Jack W Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring. Susannah Young, al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11444 2021.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140): 2020.
Samyam Rajbhandari. Conglong Li, Zhewei Yao. Minjia Zhang. Reza Yazdani Aminabadi, Am- mar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-ng Advancing mixture-of-exxt of-exeersss inference and training to power next-generation ai scale. In International conference on machine learning, pp. 13332-18342 PMLR, 2022.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. Advances in Neural Information Processing Systems, 858858838555... 2021.
Stephen Roller, Sainbayar Sukhbaatar, Jason Weston, al. Hash layers for large sparse models. advances in neural information processing systems, 74551-1566666 2021.
Andrew Rouditchenko, Saurabhchand Bhati, Edson Araujo, Samuel Thomas, Hilde Kuehne, Rogerio Feris, and James Glass. Omni-rl: Do you really need audio to fine-tune your audio Ilm? arXiv preprint 25555:2505. 094399 2025.
13

---

Published as a conference paper at ICLR 2026
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-106, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan Le Bras, and Yejin Choi. Social IQa: Common- sense reasoning about social interactions. In Kentaro Inui, Jing Jiang, Vincent Ng, and Xiaojun Wan (eds.), Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 4463 4473, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/11/D199 1444444444
URL /eattss://acllls httptthhwwlloo.ooooooooeeooeeeee 9- 1454/-
Noam Shazeer, * Azalia Mirhoseini, *Kzzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- mixture-of-experts layer. In International Conference on Learning Representations, 2017. URL https: //tptt/eee///peeeeeee... net / forum? iddiidlkkkae ckMMDqg.....
Yikang Shen, Zhen Guo. Tianle Cai, and Zengyi Qin. Jetmoe: Reaching Ilama2 performance with 0.1 m dollars. arXiv preprint arXiv:2404.00 07411,,,,,, 2024.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catan- zaro, Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.088 08055, 2019.
Luca Soldaini, Rodney Kinney, Akshita Bhagia, Dustin Sehwenk, David Atkinson, Russell Authur, Ben Bogin, Khyathi Chandu, Jennifer Dumas, Yanai Elazar, Valentin Hofmann, Ananya Jha, Sachin Kumar, Li Lucy, Xinxi Lyu, Nathan Lambert, Ian Magnusson, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew Peters, Abhilasha Ravichander, Kyle Richardson. Zejiang Shen, Emma Strubell, Nishant Subramani, Oyvind Tafjord, Evan Walsh, Luke Zettlemoyer, Noah Smith, Hannaneh Hajishirzi, Iz Beltagy, Dirk Groeneveld, Jesse Dodge, and Kyle Lo. Dolma: an open corpus of three trillion tokens for language model pretraining research. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 15725-15788, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. 18853////2222..cllloeee URL hwtps:tp////.. / ac www.anhoooooooooo.oooooooooolooe org/ soseelllonn 840/.
Kimi Team, Yifan Bai, Yiping Bao, Guanduo Chen, Jiahao Chen, Ningxin Chen, Ruijue Chen, Yanru Chen, Yuankun Chen, Yutian Chen, et al. Kimi k2: Open agentic intelligence. arXiv preprint arXiv:2507.20534, 2025.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Ilama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Ausslfarylloss-free load balancing strategy for mixture-of-experts. arXiv preprint arXiv:2408. 15664, 2024a.
Ziteng Wang, Jun Zhu, and Jianfei Chen. Remoe: Fully differentiable mixture-of-expprrs with relu routing. arXiv preprint arXiv:2412.14711, 2024b.
Geoffrey Watson. Smooth regression analysis. Sankhya: The Indian Journal of Statistics, Series A, pp. 359-372, 1964.
Tianwen Wei, Bo Zhu, Liang Zhao, Cheng Cheng, Biye Li, Weiwei Lü, Peng Cheng, Jianhao Zhang, Xiaoyu Zhang, Liang Zeng, et al. Skywork-moc: A deep dive into training techniques for mixture-of-experts language models. arXiv preprint arXiv:2406.0653 06563, 2024.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Leon Derezynski, Wei Xu, Alan Ritter, and Tim Baldwin (eds.), Proceedings of the 3rd Workshop Noisy generaede Text, pp. 94-106, Copenhagen, Denmark, September 2017. 7. Association
14

---

Published as a conference paper at ICLR 2026
for Computational Linguistics. doi: 10.18653/1//17.444133 URL UtLLttpp://cclattiooooy Exz/RR1-444335
Mengzhou Xia, Tianyu Gao, Zhiyuan Zeng, and Danqi Chen. Sheared LLaMA: Accelerating language model pre-training via structured pruning. In The Twelfth International Conference on Learning Representations, 2024. URL https: / wwtprpeeeeieeeeeee net/ forum 2 11= 100d00000....
Jing Xiong, Chengming Li, Min Yang. Xiping Hu, and Bin Hu. Expression syntax information bot- tleneck for math word problems. In Proceedings the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, pp. 2166-2171,200
Jing Xiong, Zixuan Li, Chuanyang Zheng, Zhijiang Guo, Yichun Yin, Enze Xie, Zhicheng Yang, Qingxing Cao, Haiming Wang, Xiongwei Han, et al. Dq-lore: Dual queries with low rank approximation re-ranking for in-context learning. arXiv preprint arXiv:23/0.02954 2023a.
Jing Xiong, Jianhao Shen, Ye Yuan, Haiming Wang, Yichun Yin, Zhengying Liu, Lin Li, Zhijiang
Guo, Qingxing Cao, Yinya Huang, et al. Trigo: Benchmarking formal mathematical proof reduction for generative language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 11594-11632, 2023b.
Jing Xiong. Jianghan Shen, Fanghua Ye. Chaofan Tao, Zhongwei Wan, Jianqiao Lu, Xun Wu, Chuanyang Zheng, Zhijiang Guo, Lingpeng Kong, et al. Uncomp: Uncertainty-aware long-context compressor for efficient large language model inference. 2024.
Jing Xiong, Qiujiang Chen, Fanghua Ye, Zhongwei Wan, Chuanyang Zheng, Chenyang Zhao, Hui Shen, Alexander Hanbo Li, Chaofan Tao, Haochen Tan, et al. Atts: Asynchronous test-time scaling via conformal prediction. arXiv preprint arXiv:2509.15144 2025a.
Jing Xiong, Jianghan Shen, Chuanyang Zheng, Zhongwei Wan, Chenyang Zhao, Chiwun Yang, Fanghua Ye, Hongxia Yang, Lingpeng Kong, and Ngai Wong. Parallelcomp: Parallel long-context compressor for length extrapolation. arXiv preprint 153. 2025b.
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni. Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe: An early effort on open -ixture-of-eepprrs language models. In Forty-first International Conference on Machine Learning, 2024. URL https: openreview net/ / forum? 1id= 1YDeZU8It5.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint 0093385 2025.
Donghaa Yang. Jinchuan Tian, Xu Tan, Rongjie Huang, Songxiang Liu, Xuankai Chang. Jiatong Shi, Sheng Zhao, Jiang Bian, Xixin Wu, et al. Uniaudio: audio foundation model toward universal audio generation. arXiv preprint arXhv:2310.00704 2023.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. HellaSwag: Can machine really finish your sentence? In Anna Korhonen, David Traum, and Lluís Mårquez (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/v1/119. 1472. URL https: / cl hnt ogy . org/P1 1472/.
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. arXiv preprint arXiv:2401.02388 2024.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.0 00000 2022.
Zihao Zhao, Eric Wallace, Shi Feng, Dan Klein, and Samcer Singh. Calibrate before use: Improving few-shot performance of language models. In International conference on machine learning, pp. 12697-12706. PMLR, 2021.
Shu Zhong, Mingyu Xu, Tenglong Ao, and Guang Shi. Understanding transformer from the perspec- tive of associative memory. arXiv preprint arXiv:2505. /9488, 2025.
15

---

Published as a conference paper at ICLR 2026
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoe V Le, James Laudon, et al. Mixture-of-experts with expert choice routing, Advances in Neural Information Processing Systems, 35:7103-71144 2022.
Tong Zhu, Xiaoye Qu, Daize Dong, Jiacheng Ruan, Jingqi Tong, Conghui He, and Yu Cheng. LLaMA- MoE: Building mixture-of-experts from LLaMA with continual pre-training. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen (eds.). Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing. pp. 19913-1993,, Miami, Florida, USA, November 2024. Association for Computational Linguisties. doi: 10. 11653//1/2024.mmmppp 890. URI https: / / ac antaannthhoyyl org/202 emnlp-ma in. 890/.
16

---

Published as a conference paper at ICLR 2026
A ETHICS STATEMENT
This research work is fundamentally focused on the architectural and algorithmic enhancement of the Mixtureoof-Exerrt (MoE) model paradigm. Our primary contribution involves novel integra- tion of para-mettric kernel regression methods, specifically the Nadaraya- Watson estimator, to foroouultte the gating mechanism traditionally governed by the Softmax function. This approach replaces the standard Sofsmax-based probability distribution with kernel-smoothed weighting scheme based on the KERN between an input token and each expert's 's representative vector. Conse- quently, this research does not introduce any novel, domain-specifi ethical claims societal impacts that diverge from the well-documented considcrations already associated with large-scale language models in general.
B REPRODUCIBILITY STATEMENT
comprehensive elucidation of the proposed methodology is presented in Section 3, which details the theoretical foundations and algorithmic structure of our approach. To ensure reproducibility and facilitate further research, we have made our complete code implementation publicly available. This code, which includes scripts for training, inference, and analysis, is comprehensively documented in Appendix Furthermore, the complete set hyperparameters, architectural details, and training configurations for all models discussed in our experiments are provided in Appendix D.
C
THE USE OF LLMS
For this work, mainly the Large Language Model to aid polish writing.
D MODEL CONFIGURATION
All experiments are conducted on 2 GPUs. The 125M and 350M model configuration the following.
Table Model Configurations.
125M
350M
Training sequence length
512
512
Batch size
16 x 2
16 x 2
Number of iterations
50k
50k
Dropout prob.
0.0
0.0
Attention dropout prob.
0.0
0.0
Attention head
12
16
Feature dimension
768
1024
Layer number
12
24
Optimizer
Adam
Adam
Optimizer parameter betas
[0.9, 0.95]
[0.9, 0.95]
Learning rate
6e
3e 4
Precision
float16
float16
Total Expert Number
64
64
Active Expert Number
8
8
E TIME COST AND COMPUTATIONAL COST
Theoretically, the proposed method does not have additional cost. A central advantage of the proposed gating mechanism is its computational parsimony. The operation involves calculating the L2 nomm for each expert's 's representation vector and for the input token's projection. The primary element-wise operation. Therefore, when
operation-divvsionn by the ce-norm-ccontiin
17

---

Published as a conference paper at ICLR 2026
analyzed from theoretical perspective, the proposed router introduces substantive additional time cost compared to the conventional Natmaxx-basetl approach, making it an efficient drop-in replacement.
F PERFORMANCE WITH SMALL GRANULARITY
Small Granularity Performance
4.2
Sigmaid
Validation Loss
Tanh
4.0
KERN
3.8
3.6
3.4
10K 15K 20K 25K 30K 35K 40K 45K 50K 5k
Validation Step
Figure The performance with 256 experts, 32 active experts and each expert's intermediate dimension is 96.
G THE TRAINING Loss WITH DIFFERENT METHODS
Training Loss (Length 512, Books3)
Sigmod
Tanh
Training
kxt0
10000
20000
20000
40000
50000
Training Step
Figure 8: The performance of different methods on the Books3 dataset, with active model size 125M and training length 512. There are 64 experts and 8 active experts.
The Dense model demonstrates a characteristically rapid initial learning phase, achieving a swift and substantial reduction in loss during the early stages of training. This aggressive early convergence suggests highly efficient optimization landscape for simpler, parameter-denss architectures, allowing them to quickly capitalize on low-hanging fruit within the dataset. However, this initial advantage is not sustained over the long term. As the number of training steps increases, the Dense model's loss curve begins exhibit signs of stagnation and ultimately plateaus higher value than its MoE counterparts. This pattern indicates that while the Dense model is easier to converge to reasonable solution, is ultimately constrained by its architectural limitations. The monolithie nature of its parameters appears to create lower performance ceiling, limiting its capacity to capture the complex, nuanced patterns present in the data. essence, it finds good solution quickly but lacks the expressive power to find great one. In stark contrast, the MoΣ model, particularly the one enhanced with the KERN technique, exhibits a profoundly different and more powerful learning trajectory. While its initial loss reduction may be marginally less explosive than the Dense model's demonstrates remarkable consistency and resilience throughout the entire training process. The KERN model does merely converge; it continues refine its performance, driving the loss to
18

---

Published as a conference paper at ICLR 2026
significantly lower plateau. This sustained improvement underscores superior capacity for learning
and generalization
H
THE IMPORTANCE OF CONSIDERING ALL ROUTER LOGIT
Training Loss GAP (Length 512, Books3)
-0.02
00.06
-0.06
-0.08
Losst LosSIKERNI
100000
20000
300000
40000
50000
Training Step
Figure 9: The performance of different methods the Books3 dataset, with active model size 125M and training length 512. KERN: use KERN before Top-K choice. KERN-After: use KERN after Top-K choice. The result the SAALLEEERN-Afffr-ooos(SsSSS
According to the caption in Figure 9, the result shown is the difference in loss (Loss(KERN-Afte) Loss(KERN)). Therefore, a positive value indicates that KERN has a lower loss and thus performs better. Initially, the loss difference is negative, meaning KERN-After has lower loss and performs better in early training. As training progresses, the difference becomes positive, indicating that KERN gradually achieves superior performance. This suggests that considering all router logits is important for better final performance.
I THE EFFECT OF KERN INITIALIZATION
Training Loss GAP (Length 512, Books3)
Jun www.
0.00
Training
-0.04
0.06
00.08
LosSIKERN Monto Maaalzation-LossikEEn
10000
20000
30000
40000
50000
Training Step
Figure 10: The performance of different methods on the Books3 dataset, with an active model size of 125M and training length of 512. KERN: the initial scale is 1. KERN-Monte Carlo Initialization: the Monte Carlo for the initialization, presented in Appendix L.
According to the caption in Figure 10, the result shown is the difference loss (Loss(KERN-Monte Carlo Initialization) Loss(KERNN). Therefore, positive value indicates that KERN has lower loss and thus performs better. Initially, the loss difference is negative, meaning KERN-Mone Carlo Initialization has lower loss and performs better in early training. As training progresses, the difference becomes close to zero and slightly positive, indicating that KERN gradually achieves superior performance.
19

---

Published as a conference paper at ICLR 2026
Trskiipg Lose GAP (Lensh 1024. Brcie3, d4 Act Experts)
Tahmm 644
04
-1422
2003
46666
20303
20000
20030
30083
23006
Training Stop
1aa
Figure 11: The performance of different methods on the Books3 dataset, with model parameter 520M.
J THE EFFECT OF RELU IN KERN
We compare the performance between KF.RN-without-Rel.. and KERN in Figure 11. When the ac- tive expert number is less than half of the total number, the KERN achieves comparable performance with the WERRN-wwthhut-Rell active expert number higher (such as active all experts), the KERN will achieve significantly better performance than KERN- ReL.
as the select router weights are usually positive. However, when the
Table 3: The validation loss of different methods, with training length 1024 and Books3 dataset.
Model
40K
50K
Active Expert
Softmax
4,0230
3.6848
3.5494
3.4658
3.3979
3.3447
3.2704
3.2168
3.2007
3.1714
Sigmoid
3.9648
3.6256
3.4915
3.4027
3.3347
3.2830
3.2033
3.1498
3.1327
3.1031
Tanh
3.9729
3.6457
3.5121
3.4281
3.3580
3.3034
3.2264
3.1708
3.1529
3.1224
KERN KERR--wilooutRee..
3,9398
3.6178
3.4827
3.3975
3.3278
3.2727
3.1946
3.1397
3.1214
3.0914
KERN
3,9391
3.6180
3,4842
3.3959
3.3281
3.2725
3.1954
3.1396
3.1227
3.0925
64 of 64 Active Expert
Softmax
3,9875
3.6551
3.5127
3.4252
3.3539
3.2972
3.2191
3.1641
3.1470
3.1161
Sigmoid
3.9944
3.6420
3.4991
3.4085
3.3363
3.2818
3.2035
3.1480
3.1299
3.1003
Tanh 4.0086
3.6694
3.5321
3.4489
3,3801
3.3280
3.2491
3.1910
3.1730
3.1408
KERN-wilhout-ReeLe
3.9695
3.6467
3,5045
3,4200
3.3508
3.2945
3.2139
3.1579
3.1385
3,1075
KERN 3,9513
3.6170
3,4788
3.3885
3.3199
3.2652
3.1861
3.1290
3.1081
3.0780
When the active expert number is the same as the total expert, the KERN achieves better performance than KEINN-withut-ReLL and all other methods. We compare the performance between KERN- Authout-ReLLD KERN, and other methods with all experts being active in Table 3. The KERN achieves the best performance when the expert activation ratio is relatively high, such as 64 of 64 active experts.
K PERFORMANCE WITH DIFFERENT SEEDS
Table 4: The validation loss with three random seeds, with training length 512 and Books3 dataset
Model
5K
10K
15K
20K
25K
30K
35K
40K
45K
50K
Dense
Mean
4.1536
3,9179
3.8427
3.7583
3.6899
3.6341
3.5589
3.5024
3.4737
3.4458
Variance
0.0064
0.0193
0.0078
0.0148
0.0097
0.0066
0.0299
0.0227
0.0117
0.0068
Softmax
Mean
4.3143
3,9502
3.8470
3.7472
3.6687
3.6068
3.5228
3,4569
3.4197
3.3873
Variance
0.0094
0.0196
0.0083
0.0136
0.0106
0.0086
0.0290
0.0213
0.0140
0.0050
Sigmoid
Mean
4.2605
3,9019
3.7964
3.6945
3.6153
3.5507
3.4645
3.3957
3.3571
3.3236
Variance
0.0060
0.0210
0.0078
0.0149
0.0131
0.0277
0.0222
0.0108
0.0041
Tanh
Mcan
4.2596
3,9078
3.8077
3.7063
3.6267
3.5615
3.4751
3.4058
3.3682
3.3351
Variance
0.0049
0.0173
0.0099
0.0105
0.0094
0.0111
0.0302
0.0204
0.0163
0.0050
KERN
Mean
4.2281
3,8825
3.7819
3,6823
3.6031
3.5385
3.4520
3,3828
3.3446
3.3112
Variance
0.0082
0.0194
0.0079
0.0131
0.0120
0,0085
0.0282
0.0211
0.0120
0.0041
20

---

Published as a conference paper at ICLR 2026
Table 7 presents the performance comparison of five different model architectures (Dense, Softmax, Sigmoid, Tanh, and KERN) across various training steps (from 5K to 50K) using three different random seeds. The analysis reveals several key observations: All models demonstrate consistent pattern of performance improvement as training progresses from 5K to 50K steps. The loss values decrease monotonically for all architectures, indicating successful learning convergence. The KERN model consistently achieves the best performance across all training milestones, followed closely by Sigmoid and Tanh architectures. The Dense model exhibits significantly higher variance compared to other architectures, particularly in the early training stages (0. 1699 variance at 5K steps). This suggests that the Dense architecture is more sensitive random seed initialization. In contrast, the specialized activation functions (Soflmax, Sigmoid, Tanh, and KERN) show much lower variance, indicating more stable and consistent performance across different random seeds. the final training stage (50K steps), the KERN model achieves the best performance with mean loss of 3.3112, followed by Sigmoid (3.3236) and Tanh (3.3351). The Dense model performs the worst with a mean loss of 3 3.6030, indicating that specialized activation functions provide substantial performance benefits for this task. All models show the most rapid improvement during the initial training phases (5K-20K steps), with the rate of improvement gradually slowing in later stages. This pattern suggests that while additional training continues to provide benefits, the marginal gains diminish as the models approach their performance limits on this dataset. The results demonstrate that careful selection of activation functions and normalization techniques can significantly impact model stability and final performance, with the KERN architecture emerging as the most robust and effective choice for this particular task and dataset configuration.
L IMPLEMENTATION DETAILS
In this section, we present the implementation of the proposed KERN module in PyTorch.
Trom tqdm Import tqdm.
import numpy 4s np
import
import torch.nn as nn
def (elu(x):
Implements Rectified Linear Unit (RelD) activation function.
Args:
A NumPy acray single aumerical value.
Returna:
A NumPy arzay oz single numerical value with negative values replaced by zero. wen
return mp.maxximum10 10, x)
def monte _ar__crrrrl_lllee Id, pum_eamples-100000)
samplea =
for in =qum (zange (num_ ssmples))
np.random. randr
x np. Ina.linag..ommm
relu(y)
y_sorted np np.sort (y) [::-1]
y_k = y_sorted[:k]
semples.append (y_kx=2)
return np mp.men (samples)
class KERNRousez Kodale))
def init_(
self, mhitiall_echcds"onns :otal_expart=64, top_k-8, eps-le 8,
super init_
if initial_methee =z "cre": self.acs initial
21

---

Published as a conference paper at ICLR 2026
instial_mether "sonte_carlo":
self.acele_initian mante_carle war.eaarrrr..ooe.. porrl_yy_kttee (tt__kottoo__prerrr top_k)
selt.acale on.Parameter (torch, ores (1))
self.eps
self.activation (n.
def forward x):
norm_x x.norm dims-1, keepdim-True)
"_normed = x (norn_x self.eps)
(x_normed)
return selé.scole saf.scale x_nermed
M PERFORMANCE WITH SINKHORN
Table The train loss with sinkhorn, with training length 512 and Books3 dataset
Model
5K
10K
15K
20K
25K
30K
35K
40K
45K
50K
Dense
4.1460
3,9004
3.8538
3.7480
3.6786
3.6435
3,5850
3,4900
3.4881
3,4429
Softmax
4.3198
3,9538
3.8208
3.7427
3.6518
3.5777
3,5167
3,4380
3.3966
3,3607
Sigmoid
4.2990
3.9387
3.8034
3.7263
3.6374
3.5640
3.5034
3,4233
3.3804
3.3469
KERN
4.2693
3,9241
3.7910
3.7127
3.6266
3.5525
3,4928
3,4130
3.3694
3.3345
N EXPERT LOAD BALANCE
We calculate the KI. loss for the expert load frequency and uniform distribution. The KERN achieves the lowest KL loss, suggesting that its expert load is more balanced. And the variance is also smaller.
Table 6: The KL loss for expert load frequency and uniform distribution. The KL is smaller, the load more balanced.
Method
Softmax
Sigmoid
KERN
KL Mean
0.9483
0.3162
0.2947
KL Standard Variance
0.1997
0.2207
0.1896
O
THE ODER OF NORM AND ACTIVATION FUNCTION
Table 7: The train loss with sinkhorn, with training length 512 and Books3 dataset
Model
10K
15K
20K
25K
30K
35K
40K
45K
50K
KERN (Equation 3)
4.2228
3,8683
3.7954
3.6800
3.5978
3.5527
3.4795
3.3729
3.3647
3.3132
KERN (Equation 4)
4.2165
3.8650
3.7919
3.6724
3.5910
3.5503
3.4776
3.3707
3,3608
3,3080
22


Reviews:
[{"review_id": "oa4Xeqj7i0", "reviewer": "Reviewer_dQSC", "summary": "This paper connects MoE routers to Nadaraya-Watson regression. The authors then propose KERN, a zero-cost FFN-style router using ReLU activation and ℓ2-normalization. Experiments on language modelling and NLP tasks demonstrate KERN's effectiveness.", "questions": "- In \"Advantages of the KERN router function\", why does KERN help with gradient vanishing? Sigmoid and Softmax may still have some gradient even if the input is at small value, but RELU function in KERN directly cutoff the gradient.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The authors motivates various design decisions thoroughly.\n-  This work connects kernel method perspective to MoE router design.", "weaknesses": "- Improvement is usually marginal. \n- It's a pity that the connection to Nadaraya-Watson estimator isn't well explored. It's not obvious why the router design is inspired by this perspective, they don't have much shared element, and somehow we don't see any attempt at adopting the traditional formula as a router. I think the paper will be a lot more interesting if it shows variants like Gaussian kernel.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9fVQ7LfAEe", "reviewer": "Reviewer_pDju", "summary": "This paper revisits the routing mechanism of Mixture-of-Experts (MoE) from a statistical perspective. The authors draw an analogy between the MoE router and the Nadaraya–Watson (NW) kernel regression, showing that both can be viewed as weighted aggregations over functions. Motivated by this connection, the paper introduces KERN (Kernel-Inspired Router with Normalization) — an FFN-style router that replaces Softmax with ReLU activation and ℓ₂ normalization.\nExtensive experiments across model sizes (125M–6.9B parameters), context lengths, and datasets (Arxiv, Books3, FineWeb-Edu, etc.) demonstrate consistent gains of KERN over Softmax, Sigmoid, and Tanh routers. The method adds no parameters or computational overhead, making it a practically attractive alternative for large-scale MoE systems.", "questions": "1. The paper claims a formal equivalence between softmax routing and Nadaraya–Watson kernel regression. Could the authors clarify whether this relationship is a *functional approximation* (similar weighting structure) or a *true probabilistic equivalence* under specific assumptions?  \n\n2. Under what mathematical conditions---such as temperature scaling, large-dimension limits, or specific kernel choices---can softmax routing converge to the NW estimator?  \n\n3. Given that the softmax router’s logits are learned through gradient descent rather than fixed distances, how does this affect the validity of interpreting it as a kernel regression model?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "This paper revisits the routing mechanism of Mixture-of-Experts (MoE) from a statistical perspective. The authors draw an analogy between the MoE router and the Nadaraya–Watson (NW) kernel regression, showing that both can be viewed as weighted aggregations over functions. Motivated by this connection, the paper introduces KERN (Kernel-Inspired Router with Normalization) — an FFN-style router that replaces Softmax with ReLU activation and ℓ₂ normalization.\nExtensive experiments across model sizes (125M–6.9B parameters), context lengths, and datasets (Arxiv, Books3, FineWeb-Edu, etc.) demonstrate consistent gains of KERN over Softmax, Sigmoid, and Tanh routers. The method adds no parameters or computational overhead, making it a practically attractive alternative for large-scale MoE systems.", "weaknesses": "The central theoretical claim---that Mixture-of-Experts (MoE) routing is *equivalent* to Nadaraya–Watson (NW) kernel regression---is conceptually appealing but lacks mathematical rigor. \n\nThe presented argument only shows that the softmax weighting can *approximate* a kernel-weighted aggregation under certain parameterizations, not that the two models are formally or probabilistically equivalent. \n\nIn particular, NW regression requires a symmetric and positive-definite kernel function with explicit density normalization, while the softmax router relies on asymmetric learned logits and gradient-based optimization dynamics. \n\nHence, the connection is heuristic rather than theoretically grounded. \n\nThe paper would benefit from a clear statement of the conditions (e.g., scaling limits, normalization assumptions) under which this analogy approximately holds. \n\nOverall, the claimed equivalence is oversold, weakening the theoretical soundness of the contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gYnjwX1UT4", "reviewer": "Reviewer_qEoZ", "summary": "This paper reinterprets MoE routing through the lens of the Nadaraya–Watson (NW) estimator, arguing that both FFNs and MoE can be seen as parametric NW regressors. Building on this view, the authors propose KERN (Kernel-Inspired Router with Normalization): a “FFN-style” router that replaces Softmax with a linear projection followed by L2-normalization, ReLU, and a learnable global scale, combined with standard Top-k selection. The motivation is to avoid exponential activations’ saturation, keep output scale stable, and improve expert utilization. Experiments from 520M→6.9B (active 125M→1.3B) on Books3/ArXiv and a 50B-token FineWeb-Edu pretrain report consistent loss gains and better zero-shot downstream accuracy versus Softmax/Sigmoid/Tanh routers and dense baselines.", "questions": "- How does the scale parameter $\\gamma$ evolve during training (per layer and across settings)?\n- Do gradient vanishing and saturation actually occur with Softmax and Sigmoid routers, and does KERN empirically mitigate them?\n- Does KERN improve expert load balancing?\n- How does KERN behave in MoE architectures with shared experts?", "limitations": "", "rating": 2, "confidence": 5, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "- **Clear, unified view.** The NW formulation offers a simple mathematical template that maps FFN and MoE into the same “kernel+normalization” structure, making the proposed design choices easy to state and implement.\n- **Low engineering overhead.** KERN is a drop-in change (ReLU+L2norm+γ) and the paper claims no substantive time cost relative to Softmax routing.\n- **Consistent wins in the reported setup.** Across multiple model sizes and datasets, KERN slightly but repeatedly outperforms Softmax/Sigmoid/Tanh, including at longer contexts.", "weaknesses": "- **Limited novelty.** The proposed router is essentially a straightforward substitution of the probabilistic Softmax with a normalized ReLU gate, i.e., $\\gamma \\cdot ReLU(norm(xW)))$, in replace of $Softmax(xW)$. While the Nadaraya–Watson (NW) perspective is conceptually interesting, the paper mainly **instantiates** an NW-like form without analyzing whether NW’s assumptions (kernel choice, normalization, bandwidth/scale selection) align with MoE routing objectives (load balance, capacity control, calibration). The stated motivation—“Softmax/Sigmoid suffer gradient saturation/vanishing”—is asserted rather than demonstrated (e.g., no gradient-magnitude, activation-scale, or saturation diagnostics).\n- **Insufficient comparison to closely related work.** Most components of the final design already exist in the literature, yet the paper primarily compares against vanilla Softmax/Sigmoid/Tanh routers, omitting stronger, conceptually adjacent baselines and ablations.\n\t- **ReLU routing:** \\[1\\] replaces Softmax with ReLU to obtain fully differentiable, dynamic routing.\n    - **Logit normalization:** \\[2\\] normalizes routing logits to decouple expert-centroid magnitude from assignment.\n    - **Output scaling:** Industrial implements like \\[3,4\\] adjusts the routing output scale to preserve training dynamics comparable to dense models.  \n\n- **Questionable empirical setup and reporting.** The reported gaps show Sigmoid performing markedly worse than Softmax—sometimes **larger** than the gap between Softmax-MoE and its dense counterpart—which is atypical relative to prior results \\[2,5\\] and suggests possible tuning or configuration mismatches. Critical details for effective MoE training are missing: the load-balancing objective (form and coefficient), capacity factor/dropless vs padded routing, post-Top-(k) renormalization (if any), and per-expert utilization/entropy or drop-rate statistics. Without these, it is difficult to assess stability, fairness of comparisons, or reproducibility of the claimed gains. \n\n\\[1\\] ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing https://arxiv.org/pdf/2412.14711\n\\[2\\] On the Representation Collapse of Sparse Mixture of Experts https://arxiv.org/pdf/2204.09179\n\\[3\\] DeepSeek-V3 Technical Report https://arxiv.org/abs/2412.19437\n\\[4\\] Kimi K2: Open Agentic Intelligence https://arxiv.org/abs/2507.20534\n\\[5\\] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts https://arxiv.org/pdf/2408.15664", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0UBpjdavrq", "reviewer": "Reviewer_XdAq", "summary": "This paper revisits the design of the Mixture-of-Experts (MoE) routing mechanism through the lens of Nadaraya–Watson regression, a classical kernel-based estimator. The authors identify a structural similarity between MoE routers, feedforward networks (FFNs), and kernel smoothers, and use this observation to reinterpret the MoE routing process as a form of parametric kernel regression. Building on this perspective, the authors propose a new router function called KERN (Kernel Inspired Router with Normalization), which replaces the traditional Softmax gating with a lightweight $l_2$-normalization and ReLU activation. KERN eliminates the need for projecting gating scores onto a probability simplex, simplifying training and improving numerical stability. Empirical evaluations show that KERN achieves balanced expert utilization and improved stability in large-scale MoE and LLM setups, without adding computational or parameter overhead. The approach generalizes both Softmax- and Sigmoid-based routers while maintaining the inductive biases of standard FFN layers.", "questions": "1. How does changing the order between normalization and activation (as seen in Eq. (4) vs. Eq. (3)) influence the theoretical interpretation or empirical behavior of the model? Has this difference been empirically tested or justified in the paper?\n2. Can the proposed KERN router be extended beyond MoE architectures to serve as a general routing or attention mechanism in neural networks, similar to how Softmax is used broadly across architectures?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper offers a principled reinterpretation of Mixture-of-Experts (MoE) routing through the lens of Nadaraya–Watson regression, providing theoretical clarity and unifying MoE with FFN and kernel regression models.\n2. The paper is well-organized, progressing clearly from Nadaraya–Watson regression to FFN interpretation, MoE formulation, and the final KERN methodology.\n3. The authors conduct sufficient experiments and ablation studies on multiple MoE setups, validating both the theoretical claims and practical benefits of KERN.", "weaknesses": "1. The paper appears to change the order of operations between normalization and activation layers. Specifically, Equation (4) applies the activation function before normalization, which is inconsistent with Equation (3), where normalization precedes activation. This discrepancy may affect the theoretical consistency of the proposed analogy between FFN and Nadaraya–Watson regression.\n2. The main contribution focuses on designing an alternative router function for MoE. While the proposed KERN router is elegant, its generality appears limited compared to broader architectural components like the Softmax function, which is widely used beyond MoE.\n3. The paper would benefit from a figure that clearly presents the computational flow of the proposed router function. Although the mathematical formulation is clear, a visual depiction would enhance understanding and accessibility, especially for readers less familiar with MoE internals.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)