Kernel Inspired Router with Normalization (KERN)

• The paper introduces KERN, a zero-additional-cost MoE router that replaces Softmax with linear projection, ℓ2 normalization, ReLU, and top-k masking, reframing routing as kernel regression.
• KERN leverages Nadaraya–Watson kernel regression by interpreting router outputs as kernel weights, thereby generalizing FFN-style activations for expert dispatch.
• Empirical evaluations demonstrate that KERN achieves lower validation losses and higher accuracy across benchmarks while promoting balanced expert participation.

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 (Zheng et al., 30 Sep 2025).

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 $\ell_2$0 is expert $\ell_2$1’s output and $\ell_2$2 is the router weight. With a $\ell_2$3 router,

$\ell_2$4

The paper observes that this has the same mathematical form as Nadaraya–Watson regression, whose estimator is

$\ell_2$5

Under the correspondence

$\ell_2$6

MoE is exactly

$\ell_2$7

with

$\ell_2$8

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 (Zheng et al., 30 Sep 2025).

The same paper also interprets FFNs in NW-style form. It writes

$\ell_2$9

and explicitly identifies

$\mathrm{ReLU}$0

In this view, FFN hidden units produce adaptive kernel weights, while MoE replaces static values $\mathrm{ReLU}$1 with dynamic expert outputs $\mathrm{ReLU}$2. 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

$\mathrm{ReLU}$3

with router parameters

$\mathrm{ReLU}$4

The router first computes pre-activation scores

$\mathrm{ReLU}$5

It then normalizes the entire score vector by

$\mathrm{ReLU}$6

where $\mathrm{ReLU}$7 is a small constant for numerical stability. After normalization, it applies $\mathrm{ReLU}$8 elementwise,

$\mathrm{ReLU}$9

and multiplies by a learnable global scalar

$k$0

where $k$1 is initialized to $k$2. Standard sparse MoE routing is then recovered by retaining only the top-$k$3 routed experts,

$k$4

with final weights

$k$5

The resulting MoE output is

$k$6

This is the exact KERN router described in the paper (Zheng et al., 30 Sep 2025).

An important geometric point is that KERN does not project onto the probability simplex. After $k$7, the routed weights satisfy $k$8, but generally

$k$9

The paper states explicitly that there is no additional $\mathrm{Softmax}$0 rescaling; magnitudes are controlled by the $\mathrm{Softmax}$1 normalization and the learnable scalar $\mathrm{Softmax}$2. 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}$3 with $\mathrm{Softmax}$4 normalization plus $\mathrm{Softmax}$5 does not materially increase computational complexity; and the extra operation is just norm computation, elementwise division, and $\mathrm{Softmax}$6.

3. Normalization, activation, and generalized router geometry

Normalization is central to KERN, not cosmetic. The paper gives two reasons for the $\mathrm{Softmax}$7 normalization step

$\mathrm{Softmax}$8

First, it keeps router score magnitude invariant to the number of experts $\mathrm{Softmax}$9. Second, under independently initialized experts and bounded expert outputs

$\mathrm{Sigmoid}$0

it yields stable MoE output variance at initialization: $\mathrm{Sigmoid}$1 The final step relies on the $\mathrm{Sigmoid}$2-controlled router weights. This is the paper’s main mathematical argument for scale stability (Zheng et al., 30 Sep 2025).

The choice of $\mathrm{Sigmoid}$3 is motivated in two ways. The paper argues that $\mathrm{Sigmoid}$4 and $\mathrm{Sigmoid}$5 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{Sigmoid}$6

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,

$\mathrm{Sigmoid}$7

Within this view, Softmax corresponds to an exponential activation plus implicit $\mathrm{Sigmoid}$8 normalization across experts,

$\mathrm{Sigmoid}$9

while Sigmoid-based routing reduces to

$\mathrm{Softmax}$0

without layer normalization. KERN is the FFN-style alternative that replaces exponential simplex normalization by linear scoring, vector normalization, and $\mathrm{Softmax}$1.

Empirically, the appendix compares KERN with and without $\mathrm{Softmax}$2. When few experts are active, such as $\mathrm{Softmax}$3 of $\mathrm{Softmax}$4, KERN and KERN-without-ReLU are very similar, because selected top-$\mathrm{Softmax}$5 weights are often already positive. When many experts are active, such as $\mathrm{Softmax}$6 of $\mathrm{Softmax}$7, KERN with $\mathrm{Softmax}$8 is significantly better. On Books3 at length $\mathrm{Softmax}$9 with all $\mathrm{ReLU}$0 experts active, final validation losses at $\mathrm{ReLU}$1K steps were: Softmax $\mathrm{ReLU}$2, Sigmoid $\mathrm{ReLU}$3, Tanh $\mathrm{ReLU}$4, KERN-without-ReLU $\mathrm{ReLU}$5, and KERN $\mathrm{ReLU}$6.

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}$7, initialization, and routing normalization placement (Zheng et al., 30 Sep 2025).

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 $\mathrm{ReLU}$8 and step $\mathrm{ReLU}$9K, the paper reports Dense $\ell_2$0, Softmax $\ell_2$1, Tanh $\ell_2$2, Sigmoid $\ell_2$3, and KERN $\ell_2$4. In large-scale zero-shot evaluation after training on FineWeb-Edu for $\ell_2$5K steps and $\ell_2$6B tokens, KERN exceeds Softmax at every reported scale: from $\ell_2$7 versus $\ell_2$8 at $\ell_2$9M total / $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$0M active, to $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$1 versus $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$2 at $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$3B total / $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$4B active.

The paper further reports that KERN consistently outperforms Softmax across expert counts from $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$5 to $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$6 active experts, and beats Softmax at all sparsity levels when total experts are varied from $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$7 to $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$8. With $\mathrm{MoE}(x)=\sum_{m=1}^M g_m(x)\,E_m(x),$9 experts, $\ell_2$00 active, and expert intermediate size $\ell_2$01, the Books3 result at $\ell_2$02K is Softmax $\ell_2$03, Sigmoid $\ell_2$04, Tanh $\ell_2$05, and KERN $\ell_2$06. With large sparsity at expert intermediate dimension $\ell_2$07, the Books3 result is Softmax $\ell_2$08, Sigmoid $\ell_2$09, Tanh $\ell_2$10, and KERN $\ell_2$11.

Training stability is evaluated over three seeds on Books3 length $\ell_2$12 at $\ell_2$13K steps. The reported mean losses and variances are: Softmax mean loss $\ell_2$14, variance $\ell_2$15; Sigmoid mean loss $\ell_2$16, variance $\ell_2$17; Tanh mean loss $\ell_2$18, variance $\ell_2$19; and KERN mean loss $\ell_2$20, variance $\ell_2$21. 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

$\ell_2$22

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 (Aksenov et al., 2024).

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

$\ell_2$23

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 (Landa et al., 2020). 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

$\ell_2$24

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 (Malo et al., 2018). 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

$\ell_2$25

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 (Bouhsine, 22 Feb 2026). 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$26 normalization, $\ell_2$27 activation, a global learnable scale $\ell_2$28 initialized to $\ell_2$29, and top-$\ell_2$30 masking afterward,

$\ell_2$31

It further recommends applying KERN before top-$\ell_2$32, not after. An appendix experiment compares applying KERN before top-$\ell_2$33 versus after top-$\ell_2$34, and reports that considering all router logits before top-$\ell_2$35 is important for better final performance. Early training may slightly favor post-top-$\ell_2$36, but full-logit normalization wins later (Zheng et al., 30 Sep 2025).

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-$\ell_2$37 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-$\ell_2$38 dispatch while replacing exponential normalization with norm-controlled nonnegative gating.