---
title: Kernel Inspired Router with Normalization (KERN)
url: https://www.emergentmind.com/topics/kernel-inspired-router-with-normalization-kern
type: topic
---

# Kernel Inspired Router with Normalization (KERN)

Searching arXiv for the KERN paper and closely related normalization/kernel-routing papers.
Kernel Inspired Router with Normalization (KERN) is a **zero-additional-cost** FFN-style router function for mixture-of-experts (MoE) layers that replaces the standard $\mathrm{Softmax}$ router with a linear projection, $\ell_2$ normalization, $\mathrm{ReLU}$, a learnable global scale, and standard top-$k$ masking. Its central claim is that the standard use of $\mathrm{Softmax}$ in MoE routing is mostly a convention, not a necessity derived from first principles. By reinterpreting MoE through the lens of **Nadaraya–Watson (NW) kernel regression**, KERN treats router outputs as kernel weights rather than as a mandatory projection onto the probability simplex. The paper further states that KERN generalizes both $\mathrm{Sigmoid}$- and $\mathrm{Softmax}$-based routers, and recommends the use of $\mathrm{ReLU}$ activation and $\ell_2$-normalization in the KERN router function [2509.25913].

## 1. Conceptual basis in Nadaraya–Watson regression

A standard MoE layer is written as
\[
\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),
\]
where $E_m(x)$ is expert $m$’s output and $g_m(x)$ is the router weight. With a $\mathrm{Softmax}$ router,
\[
g_m(x)=\frac{\exp(\langle w_m,x\rangle)}{\sum_{j=1}^M \exp(\langle w_j,x\rangle)}.
\]
The paper observes that this has the same mathematical form as **Nadaraya–Watson regression**, whose estimator is
\[
f(x)=\frac{\sum_{i=1}^N K(x,x_i)\,y_i}{\sum_{j=1}^N K(x,x_j)}.
\]
Under the correspondence
\[
K(x,x_i)\quad \longleftrightarrow\quad \exp(\langle w_m,x\rangle), \qquad
y_i \quad \longleftrightarrow\quad E_m(x),
\]
MoE is exactly
\[
\mathrm{MoE}(x) = \sum_{m=1}^M \frac{K(x,w_m)}{\sum_{j=1}^M K(x,w_j)}\,E_m(x),
\]
with
\[
K(x,w_m)=\exp(\langle w_m,x\rangle).
\]
This reframing is the conceptual core of KERN: routing is understood as computing **kernel-like weights** over experts, not as a special case of probabilistic classification over experts [2509.25913].

The same paper also interprets FFNs in NW-style form. It writes
\[
\mathrm{FFN}(x)=\sum_{i=1}^{h} \underbrace{\phi\!\left(\mathrm{LN}\left(\langle w_i,\Phi(x)\rangle\right)\right)}_{\text{adaptive kernel weight}} \cdot \underbrace{v_i}_{\text{value}},
\]
and explicitly identifies
\[
K(x,\{w_i,b_i\})=\phi\!\left(\langle w_i,\Phi(x)\rangle\right).
\]
In this view, FFN hidden units produce adaptive kernel weights, while MoE replaces static values $v_i$ with dynamic expert outputs $E_m(x)$. KERN therefore positions the router as an **FFN-style kernel computation** rather than as a simplex projector. This suggests that the MoE router can be redesigned by borrowing from FFN practice without changing the surrounding MoE formulation.

## 2. Formal definition of the KERN router

The paper defines KERN on a router input representation
\[
\Phi(x)\in\mathbb{R}^d
\]
with router parameters
\[
W_s\in\mathbb{R}^{M\times d},\qquad b_s\in\mathbb{R}^M.
\]
The router first computes pre-activation scores
\[
s(x)=W_s\Phi(x)+b_s.
\]
It then normalizes the entire score vector by
\[
\bar{s}(x)=\frac{s(x)}{\|s(x)\|_2+\varepsilon},
\]
where $\varepsilon>0$ is a small constant for numerical stability. After normalization, it applies $\mathrm{ReLU}$ elementwise,
\[
r(x)=\mathrm{ReLU}(\bar{s}(x)),
\]
and multiplies by a learnable global scalar
\[
\hat{g}(x)=\gamma\, r(x),
\]
where $\gamma$ is initialized to $1$. Standard sparse MoE routing is then recovered by retaining only the top-$k$ routed experts,
\[
\mathcal{T}_k(x)=\operatorname{TopKIndices}\!\big(\hat{g}(x),k\big),
\]
with final weights
\[
g_m(x)=\hat{g}_m(x)\,\mathbf{1}[m\in \mathcal{T}_k(x)].
\]
The resulting MoE output is
\[
\mathrm{MoE}_{\mathrm{KERN}}(x)=\sum_{m=1}^M g_m(x)\,E_m(x).
\]
This is the exact KERN router described in the paper [2509.25913].

An important geometric point is that KERN does **not** project onto the probability simplex. After $\mathrm{ReLU}$, the routed weights satisfy $g_m(x)\ge 0$, but generally
\[
\sum_{m=1}^M g_m(x)\neq 1.
\]
The paper states explicitly that there is **no additional $\ell_1$ rescaling**; magnitudes are controlled by the $\ell_2$ normalization and the learnable scalar $\gamma$. Accordingly, KERN lives in a nonnegative, norm-controlled space rather than the simplex.

The method is described as **zero-additional-cost** because it changes only the router function. It introduces **no new experts** and **no extra expert-side parameters**; the router still uses the same linear projection shape as standard MoE routers; replacing $\mathrm{Softmax}$ with $\ell_2$ normalization plus $\mathrm{ReLU}$ does not materially increase computational complexity; and the extra operation is just norm computation, elementwise division, and $\mathrm{ReLU}$.

## 3. Normalization, activation, and generalized router geometry

Normalization is central to KERN, not cosmetic. The paper gives two reasons for the $\ell_2$ normalization step
\[
\bar{s}(x)=\frac{s(x)}{\|s(x)\|_2+\varepsilon}.
\]
First, it keeps router score magnitude invariant to the number of experts $M$. Second, under independently initialized experts and bounded expert outputs
\[
\|E_m(x)\|_2=\mathcal{O}(1),
\]
it yields stable MoE output variance at initialization:
\[
\mathbb{E}\left[\left\|\mathrm{MoE}_{\mathrm{KERN}}(x)\right\|_2^2\right]
=
\mathbb{E}\left[\left\|\sum_{m=1}^M g_m(x)E_m(x)\right\|_2^2\right]
=
\sum_{m=1}^M g_m(x)^2\,\mathbb{E}\left[\|E_m(x)\|_2^2\right]
=
\mathcal{O}(1)\cdot \sum_{m=1}^M g_m(x)^2
=
\mathcal{O}(1).
\]
The final step relies on the $\ell_2$-controlled router weights. This is the paper’s main mathematical argument for scale stability [2509.25913].

The choice of $\mathrm{ReLU}$ is motivated in two ways. The paper argues that $\mathrm{Softmax}$ and $\mathrm{Sigmoid}$ depend on exponentials and can be overly sensitive to input magnitude, causing sharp value explosion, saturation, near-zero weights for some experts, and vanishing gradients for underused experts. By contrast,
\[
\mathrm{ReLU}(z)=\max(0,z)
\]
is piecewise linear and does not saturate on the positive side. The second argument is architectural: modern FFNs commonly use ReLU-like activations or linear outputs rather than exponentials, so KERN mirrors standard FFN design more closely than Softmax routing does.

The paper also embeds Softmax and Sigmoid in a generalized FFN-style routing family,
\[
g_m(x)=\phi\!\left(\mathrm{LN}\left(\langle w_m,\Phi(x)\rangle\right)\right).
\]
Within this view, Softmax corresponds to an exponential activation plus implicit $\ell_1$ normalization across experts,
\[
g_m(x)=\frac{\exp(\langle w_m,\Phi(x)\rangle)}{\sum_j \exp(\langle w_j,\Phi(x)\rangle)},
\]
while Sigmoid-based routing reduces to
\[
g_m(x)=\sigma(\langle w_m,\Phi(x)\rangle)
\]
without layer normalization. KERN is the FFN-style alternative that replaces exponential simplex normalization by linear scoring, vector normalization, and $\mathrm{ReLU}$.

Empirically, the appendix compares KERN with and without $\mathrm{ReLU}$. When few experts are active, such as $8$ of $64$, KERN and KERN-without-ReLU are very similar, because selected top-$k$ weights are often already positive. When many experts are active, such as $64$ of $64$, KERN with $\mathrm{ReLU}$ is significantly better. On Books3 at length $1024$ with all $64$ experts active, final validation losses at $50$K steps were: Softmax $3.1161$, Sigmoid $3.1003$, Tanh $3.1408$, KERN-without-ReLU $3.1075$, and KERN $3.0780$.

## 4. Empirical evaluation

The experiments cover language modeling on **Arxiv** and **Books3**, large-scale pretraining on **FineWeb-Edu**, downstream zero-shot evaluation, different model scales, different sequence lengths, different expert granularities, different sparsity levels, and ablations on $\mathrm{ReLU}$, initialization, and routing normalization placement [2509.25913].

A representative subset of the reported comparisons is shown below.

| Setting | Softmax | KERN |
|---|---:|---:|
| Arxiv, length 512, active 125M / total 520M, step 50K | 1.8781 | **1.8291** |
| Books3, length 512, step 50K | 3.3882 | **3.3080** |
| Books3, length 1024, step 50K | 3.1714 | **3.0914** |
| Books3, length 2048, step 50K | 3.0442 | **2.9535** |
| FineWeb-Edu, 520M total / 125M active, average accuracy | 49.88 | **52.14** |
| FineWeb-Edu, 1.7B total / 350M active, average accuracy | 52.46 | **55.13** |
| FineWeb-Edu, 6.9B total / 1.3B active, average accuracy | 56.49 | **58.88** |

These results are accompanied by comparisons against Dense, Sigmoid, and Tanh routers. On Books3 at length $2048$ and step $50$K, the paper reports Dense $3.1249$, Softmax $3.0442$, Tanh $2.9868$, Sigmoid $2.9635$, and KERN $2.9535$. In large-scale zero-shot evaluation after training on FineWeb-Edu for $50$K steps and $50$B tokens, KERN exceeds Softmax at every reported scale: from $52.14$ versus $49.88$ at $520$M total / $125$M active, to $58.88$ versus $56.49$ at $6.9$B total / $1.3$B active.

The paper further reports that KERN consistently outperforms Softmax across expert counts from $4$ to $32$ active experts, and beats Softmax at all sparsity levels when total experts are varied from $32$ to $256$. With $256$ experts, $8$ active, and expert intermediate size $96$, the Books3 result at $50$K is Softmax $3.4150$, Sigmoid $3.3302$, Tanh $3.3276$, and KERN $3.2962$. With large sparsity at expert intermediate dimension $384$, the Books3 result is Softmax $3.3761$, Sigmoid $3.2760$, Tanh $3.2972$, and KERN $3.2672$.

Training stability is evaluated over three seeds on Books3 length $512$ at $50$K steps. The reported mean losses and variances are: Softmax mean loss $3.3873$, variance $0.0050$; Sigmoid mean loss $3.3236$, variance $0.0041$; Tanh mean loss $3.3351$, variance $0.0050$; and KERN mean loss $3.3112$, variance $0.0041$. The paper interprets this as KERN achieving the best mean and low variance.

## 5. Relation to adjacent kernel-and-normalization lines of work

KERN belongs to a broader line of work in which kernel geometry and normalization are treated as co-design variables rather than as independent modules. A closely related example is the linear-attention model "ReBased," which introduces a learnable second-order polynomial kernel with normalization applied before kernel evaluation. Its final ablation form is
\[
(\gamma\cdot norm(x) + \beta)^2,
\]
and the paper explicitly states that normalization is applied to queries and keys **before kernel evaluation**, not as a separate output-side normalizer. It reports that this design improves long-context retrieval and training stability relative to the fixed Based kernel, which is directly relevant to KERN’s emphasis on normalized pre-kernel feature geometry [2402.10644].

A second relevant normalization principle comes from the study of Gaussian-kernel affinity matrices under heteroskedastic noise. That work shows that heteroskedastic noise induces a **left-right diagonal scaling bias**
\[
\widetilde K \approx \operatorname{diag}(a)\, K \,\operatorname{diag}(a),
\]
and that row-stochastic normalization does not remove this two-sided distortion, whereas doubly-stochastic normalization does. For a KERN-style design, the explicit transfer made in the paper summary is that if routing scores are built from Gaussian similarities between noisy embeddings, then enforcing a symmetric doubly-stochastic scaling can remove a class of multiplicative per-node biases that row-normalized soft assignments do not remove [2006.00402]. This suggests a broader normalization lesson: not all normalizations remove the same kernel distortions.

A third line of work derives effective normalization kernels of the form
\[
H = \mathbb{D}_{\left(\frac{k}{x}\right)}\, W\, \mathbb{D}_{\left(\frac{k}{b}\right)},
\]
showing that an effective interaction kernel may require both source-side and target-side diagonal reweighting around a base kernel. In that setting, a symmetric lower-level interaction kernel can induce an effective **non-symmetric** normalization kernel [1804.05964]. A plausible implication is that KERN-like routers need not be limited to plain similarity matrices or simplex projections; pre- and post-scaling around a base interaction may be structurally useful when routing is normalization-aware.

A broader related trend places normalization inside the kernel score itself. The yat-product kernel
\[
K_\E(\mathbf{x},\mathbf{a}) = \frac{(\mathbf{x}^\top \mathbf{a})^2}{\|\mathbf{x}-\mathbf{a}\|^2+\varepsilon}
\]
was used as an attention score and as the sole non-linearity in Neural Matter Networks, with the paper arguing that normalization shifts into the kernel itself via the denominator rather than relying on separate normalization layers [2603.12276]. This is not KERN, but it exemplifies the same general movement away from treating routing score computation and normalization as separable design stages.

## 6. Practical recommendations, caveats, and unresolved questions

The paper’s practical recommendation is explicit: use a **linear router projection**, **$\ell_2$ normalization**, **$\mathrm{ReLU}$ activation**, a **global learnable scale $\gamma$** initialized to $1$, and **top-$k$ masking afterward**,
\[
s=W_s\Phi(x)+b_s,\qquad
\bar{s}=\frac{s}{\|s\|_2+\varepsilon},\qquad
\hat{g}=\gamma\,\mathrm{ReLU}(\bar{s}).
\]
It further recommends applying KERN **before** top-$k$, not after. An appendix experiment compares applying KERN before top-$k$ versus after top-$k$, and reports that considering **all router logits before top-$k$** is important for better final performance. Early training may slightly favor post-top-$k$, but full-logit normalization wins later [2509.25913].

At the same time, the paper states several caveats. It does **not** provide a formal convergence theorem for KERN. It does **not** deeply analyze auxiliary load balancing objectives. It mainly validates the method on language modeling and LLM pretraining setups, so broader modality coverage is not directly tested there. In very sparse top-$k$ settings, KERN-without-ReLU can be very close to KERN, so ReLU’s gain may depend on activation density. The paper also claims that KERN promotes more balanced expert participation and better utilization, primarily through reduced saturation, normalized score magnitudes, and ReLU sparsity, but it does **not** present a dedicated mathematical load-balancing theorem or a detailed auxiliary-loss redesign.

Within those limits, KERN is presented as a drop-in replacement for Softmax routing in MoE. Its distinctive features are that it is motivated by the NW interpretation of routing, that it mirrors FFN-style computation rather than simplex projection, and that it relies on explicit vector normalization before sparse expert selection. The result is a router that preserves standard top-$k$ dispatch while replacing exponential normalization with norm-controlled nonnegative gating.

Source: https://www.emergentmind.com/topics/kernel-inspired-router-with-normalization-kern