% !TEX program = pdflatex
% !BIB program = bibtex
% ===========================================================================
%  PubArx LaTeX template -- example paper
%
%  This file demonstrates every feature of the pubarx class. Keep the
%  structure and replace the example content with your own. To compile:
%
%      pdflatex main
%      bibtex main
%      pdflatex main
%      pdflatex main
%
%  or upload the folder to Overleaf (Menu > Compiler > pdfLaTeX).
%  The results reported below are illustrative only.
%  Every option and command is documented in README.md.
% ===========================================================================
\documentclass[a4paper]{pubarx}
% Options (combine with commas), for example \documentclass[letterpaper,review]{pubarx}:
%   letterpaper   US Letter paper instead of A4
%   onecolumn     a single column (11 pt) instead of two columns (10 pt)
%   review        line numbers and wider line spacing for reviewers
%   columnrule    a thin rule between the two columns
%   nologo        no PubArx band, e.g. for a journal submission
%   font=newtx    Times-like fonts instead of Libertinus (or font=lm)
%   notheorems    do not predefine theorem environments

% The class already loads amsmath, amsthm, graphicx, xcolor, booktabs,
% caption, enumitem, natbib and hyperref: do not load them again with
% options (use \hypersetup{...} or \captionsetup{...} instead).
% Packages used by this example only:
\usepackage{tikz}
\usetikzlibrary{arrows.meta}

\DeclareMathOperator*{\argmax}{arg\,max}
\newcommand{\norm}[1]{\lVert #1 \rVert}
\newcommand{\R}{\mathbb{R}}

\begin{document}

% --- Front matter ----------------------------------------------------------
% Everything up to \maketitle builds the full-width title block on page 1 and
% fills in the PDF metadata (title, authors, keywords) that indexers read.

\title[Stable routing in sparse mixture-of-experts models]{Stable Token Routing
  in Sparse Mixture-of-Experts Models via Temperature-Controlled
  Gating\thanks{An example paper produced with the PubArx \LaTeX{} template.
  Its numerical results are illustrative.}}

% One \author per author, in order. The optional argument lists affiliation
% labels; follow it with \orcid, \corresponding and \equalcontrib as needed.
\author[1]{Maryam Rahimi}\orcid{0000-0002-1825-0097}\equalcontrib
  \corresponding{m.rahimi@northbridge.example}
\author[1,2]{Daniel K. Okafor}\orcid{0000-0000-1234-5672}\equalcontrib
\author[2]{Lena Fischer}\orcid{0000-0000-7654-3213}

\affiliation[1]{Department of Computer Science, University of Northbridge,
  Toronto, Canada}
\affiliation[2]{Institute for Data Science, Halden Technical University,
  Oslo, Norway}

\runningauthor{Rahimi, Okafor and Fischer}

% PubArx band. Leave \pubarxid out until PubArx has assigned your identifier
% (it is shown on your paper page), then add it to later versions;
% \pubarxversion has no effect without it.
\pubarxid{2026.00000}
\pubarxversion{1}
\pubarxstatus{Preprint -- not peer reviewed}
\pubarxlicense{CC BY 4.0}
\pubarxsubject{Machine Learning}

\received{14 September 2026}
\posted{2 October 2026}

\keywords{mixture-of-experts, sparse routing, Lipschitz continuity,
  load balancing, language models}

\begin{abstract}
Sparse mixture-of-experts (MoE) layers scale neural networks by sending each
token to a few expert sub-networks chosen by a learned router. Routers are
brittle: a small change to a token representation can flip its expert
assignment, which destabilises training and wastes capacity. We study
temperature-controlled softmax gating and prove that the gate is Lipschitz
continuous with constant $\norm{W}_2/(2\tau)$, where $W$ holds the router
weights and $\tau$ is the temperature. The bound yields an explicit radius
within which top-1 assignments cannot change. Guided by this result, we
propose a temperature schedule that keeps the router's Lipschitz constant
fixed during training. On a language model with 125 million active
parameters, the schedule lowers validation perplexity from 14.5 to 13.5 and
reduces token dropping fourfold compared with Switch-style routing, at no
additional cost. Our analysis suggests that the router temperature, not only
auxiliary losses, should be treated as a first-class stability parameter.
\end{abstract}

\begin{highlights}
  \item Softmax routers are Lipschitz with constant $\norm{W}_2/(2\tau)$.
  \item The bound certifies a radius within which top-1 routing cannot change.
  \item A temperature schedule cuts token dropping fourfold at no extra cost.
\end{highlights}

\maketitle

% --- Main text -------------------------------------------------------------

\section{Introduction}\label{sec:intro}

Mixture-of-experts models combine several specialised networks with a gate
that decides which of them should handle each input~\citep{jacobs1991adaptive}.
\citet{shazeer2017outrageously} made the idea practical for deep
learning~\citep{lecun2015deep} by activating only a few experts per token, and
sparse expert layers now power some of the largest transformer language
models~\citep{vaswani2017attention,lepikhin2021gshard,fedus2022switch}.

Sparsity comes at a price. Because each token is processed only by the experts
with the largest gate values, the routing decision is discontinuous: two
nearly identical tokens may be sent to different experts. Practitioners
counter the resulting instability with auxiliary load-balancing losses,
capacity limits and careful initialisation, but the role of the router
temperature has received little attention. This paper makes three
contributions:
\begin{itemize}
  \item We prove that temperature-controlled softmax gating is Lipschitz
    continuous with constant $\norm{W}_2/(2\tau)$
    (Theorem~\ref{thm:lipschitz}).
  \item We derive a certified radius within which the top-1 expert of a
    token cannot change (Corollary~\ref{cor:radius}).
  \item We propose a temperature schedule that holds the Lipschitz constant
    fixed, and show that it improves both perplexity and expert utilisation
    (Section~\ref{sec:experiments}).
\end{itemize}

\section{Background}\label{sec:background}

Let $x \in \R^d$ be a token representation and let $E_1, \dots, E_E$ be the
expert networks of a layer. A linear router with weights
$W \in \R^{E \times d}$ and temperature $\tau > 0$ produces the gate
\begin{equation}
  g_\tau(x) = \operatorname{softmax}\bigl(Wx/\tau\bigr),
  \label{eq:gate}
\end{equation}
where $\operatorname{softmax}(z)_i = \exp(z_i) / \sum_{j=1}^{E} \exp(z_j)$.
The temperature controls how decisive the router is: as $\tau \to 0$ the gate
approaches a one-hot vector, and as $\tau \to \infty$ it approaches the
uniform distribution.

\begin{definition}[Top-$k$ routing]\label{def:topk}
For $k \in \{1, \dots, E\}$, let $\mathcal{T}_k(x)$ be the indices of the $k$
largest entries of $g_\tau(x)$, with ties broken by index. A top-$k$
mixture-of-experts layer computes
\begin{equation}
  y = \sum_{i \in \mathcal{T}_k(x)} g_{\tau,i}(x)\, E_i(x).
  \label{eq:moe}
\end{equation}
\end{definition}

Figure~\ref{fig:router} illustrates top-2 routing. Only the selected experts
are evaluated, so the cost of a layer grows with $k$ rather than with $E$.

\begin{figure}[t]
  \centering
  \begin{tikzpicture}[
      font=\footnotesize,
      block/.style={draw=pubarxnavy, line width=0.6pt, rounded corners=2pt,
        minimum height=6.5mm, inner sep=3pt, align=center},
      expert/.style={block, minimum width=11mm, fill=pubarxtint},
      chosen/.style={block, minimum width=11mm, draw=pubarxaccent, fill=pubarxwarm},
      flow/.style={-{Stealth[length=1.6mm]}, line width=0.6pt, draw=pubarxnavy!80},
      idle/.style={-{Stealth[length=1.4mm]}, line width=0.4pt, draw=pubarxnavy!35, dashed},
      route/.style={-{Stealth[length=1.8mm]}, line width=1.1pt, draw=pubarxaccent}]
    \node[block, minimum width=9mm] (x) at (0,0) {$x_t$};
    \node[block, minimum width=13mm, fill=pubarxnavy, text=white] (r) at (1.75,0)
      {router\\[-1pt]$g_\tau$};
    \node[expert] (e1) at (3.95,1.35) {$E_1$};
    \node[chosen] (e2) at (3.95,0.45) {$E_2$};
    \node[expert] (e3) at (3.95,-0.45) {$E_3$};
    \node[chosen] (e4) at (3.95,-1.35) {$E_4$};
    \node[draw=pubarxnavy, circle, line width=0.6pt, inner sep=1pt,
      minimum size=5.5mm] (sum) at (5.75,0) {$\Sigma$};
    \node[block, minimum width=9mm] (y) at (7.15,0) {$y_t$};
    \draw[flow] (x) -- (r);
    \draw[idle] (r.east) -- (e1.west);
    \draw[idle] (r.east) -- (e3.west);
    \draw[route] (r.east) -- node[above, sloped, inner sep=1.5pt] {$0.62$} (e2.west);
    \draw[route] (r.east) -- node[below, sloped, inner sep=1.5pt] {$0.31$} (e4.west);
    \draw[flow] (e2.east) -- (sum);
    \draw[flow] (e4.east) -- (sum);
    \draw[flow] (sum) -- (y);
  \end{tikzpicture}
  \caption{Top-2 routing of a token $x_t$. The router scores every expert, but
    only the two with the largest gate values (terracotta) process the token;
    their outputs are summed with the gate values as weights,
    as in~\eqref{eq:moe}.}
  \label{fig:router}
\end{figure}

\section{Lipschitz routing}\label{sec:method}

Our analysis rests on a simple fact about the softmax Jacobian.
\citet{gao2017properties} show that the softmax function with inverse
temperature $\lambda$ is $\lambda$-Lipschitz. The following result sharpens
this by a factor of two and makes the dependence on the router weights
explicit; $\norm{W}_2$ denotes the spectral norm.

\begin{theorem}[Lipschitz routing]\label{thm:lipschitz}
For all $x, y \in \R^d$,
\begin{equation}
  \norm{g_\tau(x) - g_\tau(y)}_2 \le \frac{\norm{W}_2}{2\tau}\, \norm{x - y}_2.
  \label{eq:lipschitz}
\end{equation}
\end{theorem}

\begin{proof}
Write $\sigma$ for the softmax function and $p = \sigma(z)$. Its Jacobian is
$J(z) = \operatorname{diag}(p) - pp^\top$, a symmetric positive semidefinite
matrix. By Gershgorin's circle theorem~\citep{horn2012matrix}, every eigenvalue
of $J(z)$ lies within $r_i = \sum_{j \ne i} p_i p_j = p_i(1 - p_i)$ of some
diagonal entry $J_{ii} = p_i(1 - p_i)$. Hence every eigenvalue is at most
$2p_i(1 - p_i) \le 1/2$, and $\norm{J(z)}_2 \le 1/2$. With
$z(t) = W\bigl(y + t(x - y)\bigr)/\tau$, the fundamental theorem of calculus
gives
\[
  g_\tau(x) - g_\tau(y) = \int_0^1 J\bigl(z(t)\bigr)\, \frac{W}{\tau}\, (x - y)
  \,\mathrm{d}t,
\]
and taking norms yields~\eqref{eq:lipschitz}.
\end{proof}

The bound translates directly into a guarantee for individual tokens.

\begin{corollary}[Certified routing radius]\label{cor:radius}
Let $m(x) > 0$ be the gap between the largest and the second-largest entry of
$g_\tau(x)$. Every $y$ with
\[
  \norm{x - y}_2 < \frac{\sqrt{2}\, \tau\, m(x)}{\norm{W}_2}
\]
is routed to the same top-1 expert as $x$.
\end{corollary}

The proof is given in Appendix~\ref{app:proof}.

\begin{remark}\label{rem:temperature}
Theorem~\ref{thm:lipschitz} makes the trade-off explicit. Halving $\tau$
sharpens routing but doubles the Lipschitz constant, so small perturbations,
such as those caused by noisy gradients early in training, are more likely to
change expert assignments.
\end{remark}

\section{A temperature schedule}\label{sec:schedule}

During training the spectral norm of $W$ typically grows, which silently
increases the Lipschitz constant in~\eqref{eq:lipschitz}. We therefore keep
the constant at a target value $L_\star$ by adapting the temperature. The
spectral norm is estimated with one step of power iteration per update, as in
spectral normalisation~\citep{miyato2018spectral}:
\begin{align}
  v_{t+1} &= \frac{W_t^\top W_t v_t}{\norm{W_t^\top W_t v_t}_2}, \qquad
  \hat\sigma_t = \norm{W_t v_t}_2, \label{eq:power}\\
  \tau_t &= \max\Bigl\{\tau_{\min},\; \frac{\hat\sigma_t}{2L_\star}\Bigr\}.
  \label{eq:schedule}
\end{align}
The model is trained with the usual auxiliary load-balancing
term~\citep{fedus2022switch},
\begin{equation}
  \mathcal{L} = \mathcal{L}_{\mathrm{LM}} + \alpha E \sum_{i=1}^{E} f_i P_i,
  \label{eq:loss}
\end{equation}
where $f_i$ is the fraction of tokens whose top-1 expert is~$i$ and $P_i$ is
the mean gate value of expert~$i$ in the batch. The schedule adds one
matrix--vector product per layer and step, which is negligible.

\section{Experiments}\label{sec:experiments}

We train decoder-only transformer language models~\citep{vaswani2017attention}
with 12 layers and width 768, replacing every other feed-forward block with an
MoE layer of $E = 8$ experts and top-2 routing. All models have about
125 million active parameters per token and are trained for 60{,}000 steps
with Adam~\citep{kingma2015adam}; Appendix~\ref{app:details} lists the
complete configuration.

Table~\ref{tab:results} compares the schedule~\eqref{eq:schedule} with
Switch-style top-1 routing~\citep{fedus2022switch} and with fixed
temperatures. The schedule
reaches the lowest perplexity, while a high fixed temperature balances load
slightly better but routes less decisively.\footnote{Load CV is the
coefficient of variation of the number of tokens per expert; 0 means a
perfectly even load.} Figure~\ref{fig:training} shows that the schedule also
converges faster and spreads tokens almost uniformly across experts.

\begin{table}[t]
  \caption{Validation perplexity (PPL), load imbalance (CV) and share of
    dropped tokens after 60{,}000 steps, for models with the same number of
    active parameters (illustrative values). Lower is better; the best value
    in each column is in bold.}
  \label{tab:results}
  \centering\small
  \setlength{\tabcolsep}{4.5pt}
  \begin{tabular}{@{}lcccc@{}}
    \toprule
    Model & Params & PPL & Load CV & Drop (\%) \\
    \midrule
    Dense & 125M & 16.9 & -- & -- \\
    Switch, top-1~\citep{fedus2022switch} & 125M & 14.5 & 0.41 & 3.8 \\
    Fixed $\tau = 0.5$ & 125M & 15.9 & 0.58 & 6.1 \\
    Fixed $\tau = 2.0$ & 125M & 13.9 & \textbf{0.12} & \textbf{0.4} \\
    Scheduled (ours) & 125M & \textbf{13.5} & 0.15 & 0.9 \\
    \bottomrule
  \end{tabular}
\end{table}

\begin{figure*}[t]
  \centering
  \begin{minipage}[b]{0.49\textwidth}
    \centering
    \resizebox{\linewidth}{!}{%
    \begin{tikzpicture}[font=\footnotesize]
      \foreach \y in {1,2,3,4} {
        \draw[pubarxrule, line width=0.3pt] (0,\y) -- (7.2,\y);
      }
      \draw[pubarxmuted, line width=0.5pt] (0,0) -- (7.2,0);
      \foreach \x/\lab in {0/0, 1.2/10, 2.4/20, 3.6/30, 4.8/40, 6.0/50, 7.2/60} {
        \draw[pubarxmuted, line width=0.5pt] (\x,0) -- (\x,-0.08);
        \node[below, text=pubarxmuted] at (\x,-0.08) {\lab};
      }
      \foreach \y/\lab in {0/10, 1/15, 2/20, 3/25, 4/30} {
        \node[left, text=pubarxmuted] at (0,\y) {\lab};
      }
      \node[text=pubarxmuted] at (3.6,-0.8) {training steps (thousands)};
      \node[text=pubarxmuted, rotate=90] at (-0.9,2) {validation perplexity};
      \draw[pubarxnavy!45, line width=1pt, dashed] plot coordinates
        {(0.6,3.30) (1.2,2.40) (2.4,1.64) (3.6,1.26) (4.8,1.06) (6.0,0.96) (7.2,0.90)};
      \draw[pubarxmuted, line width=1pt, dotted] plot coordinates
        {(0.6,3.40) (1.2,2.60) (2.4,1.80) (3.6,1.48) (4.8,1.32) (6.0,1.22) (7.2,1.18)};
      \draw[pubarxaccent, line width=1.4pt] plot coordinates
        {(0.6,3.20) (1.2,2.20) (2.4,1.40) (3.6,1.04) (4.8,0.86) (6.0,0.76) (7.2,0.70)};
      \begin{scope}[shift={(4.25,3.78)}]
        \draw[pubarxaccent, line width=1.4pt] (0,0) -- (0.5,0);
        \node[right] at (0.55,0) {scheduled (ours)};
        \draw[pubarxnavy!45, line width=1pt, dashed] (0,-0.3) -- (0.5,-0.3);
        \node[right] at (0.55,-0.3) {Switch, top-1};
        \draw[pubarxmuted, line width=1pt, dotted] (0,-0.6) -- (0.5,-0.6);
        \node[right] at (0.55,-0.6) {fixed $\tau = 0.5$};
      \end{scope}
      \node[font=\sffamily\bfseries\small, text=pubarxnavy] at (-0.9,4.35) {a};
    \end{tikzpicture}}
  \end{minipage}\hfill
  \begin{minipage}[b]{0.49\textwidth}
    \centering
    \resizebox{\linewidth}{!}{%
    \begin{tikzpicture}[font=\footnotesize]
      \foreach \y in {0.8,1.6,2.4,3.2,4.0} {
        \draw[pubarxrule, line width=0.3pt] (0,\y) -- (7.4,\y);
      }
      \foreach \y/\lab in {0/0, 0.8/5, 1.6/10, 2.4/15, 3.2/20, 4.0/25} {
        \node[left, text=pubarxmuted] at (0,\y) {\lab};
      }
      \foreach \e/\base/\ours in {1/3.52/2.18, 2/2.88/2.10, 3/2.40/2.05,
          4/1.92/2.00, 5/1.76/1.98, 6/1.44/1.94, 7/1.12/1.90, 8/0.96/1.86} {
        \pgfmathsetmacro{\cx}{0.55+(\e-1)*0.9}
        \fill[pubarxnavy!35] (\cx-0.34,0) rectangle (\cx-0.02,\base);
        \fill[pubarxaccent] (\cx+0.02,0) rectangle (\cx+0.34,\ours);
        \node[below, text=pubarxmuted] at (\cx,0) {$E_{\e}$};
      }
      \draw[pubarxmuted, line width=0.5pt] (0,0) -- (7.4,0);
      \draw[pubarxnavy, line width=0.6pt, dash pattern=on 3pt off 2pt] (0,2.0) -- (7.4,2.0);
      \node[above left, text=pubarxnavy, font=\scriptsize] at (7.4,2.0) {uniform load, 12.5\%};
      \node[text=pubarxmuted] at (3.7,-0.8) {expert};
      \node[text=pubarxmuted, rotate=90] at (-0.9,2) {share of tokens (\%)};
      \begin{scope}[shift={(4.25,3.70)}]
        \fill[pubarxnavy!35] (0,-0.09) rectangle (0.3,0.09);
        \node[right] at (0.35,0) {Switch, top-1};
        \fill[pubarxaccent] (0,-0.39) rectangle (0.3,-0.21);
        \node[right] at (0.35,-0.3) {scheduled (ours)};
      \end{scope}
      \node[font=\sffamily\bfseries\small, text=pubarxnavy] at (-0.9,4.35) {b};
    \end{tikzpicture}}
  \end{minipage}
  \caption{Training dynamics and expert load (illustrative values).
    (a)~Validation perplexity during training: the temperature schedule
    converges faster, and to a lower value, than Switch-style top-1 routing
    and a fixed low temperature. (b)~Share of tokens processed by each expert
    after training: the schedule keeps every expert close to the uniform load
    of 12.5\%.}
  \label{fig:training}
\end{figure*}

\section{Discussion and limitations}\label{sec:discussion}

The certified radius of Corollary~\ref{cor:radius} is a worst-case guarantee.
In our runs most tokens sit well inside it, which suggests that a tighter,
data-dependent analysis is possible. Three limitations deserve mention:
\begin{enumerate}
  \item The bound concerns a single linear router. Routers with hidden
    layers or noise require a separate analysis.
  \item Power iteration underestimates the spectral norm early in training,
    so the schedule is conservative only after a few hundred steps.
  \item We evaluated models with up to 125 million active parameters; the
    behaviour at larger scales remains to be confirmed.
\end{enumerate}

\section{Conclusion}\label{sec:conclusion}

Temperature is the most direct handle on the stability of a sparse router.
By keeping the Lipschitz constant of the gate fixed, a one-line schedule
improves perplexity and expert utilisation without extra computation. We hope
that this view encourages a more principled treatment of routing in sparse
models.

% --- Back matter -----------------------------------------------------------

\section*{Acknowledgements}
We thank the members of the Northbridge Machine Learning Group for helpful
discussions, and the PubArx community for their open reviews of an earlier
version of this preprint.

\section*{Data availability}
Code, configuration files and training logs are available at
\url{https://example.org/pubarx-moe-routing} under the MIT licence, following
the FAIR principles~\citep{wilkinson2016fair}.

\section*{Author contributions}
Contributions are described with the CRediT taxonomy~\citep{brand2015beyond}.
\textbf{M.~Rahimi:} conceptualization, formal analysis, writing -- original
draft. \textbf{D.~K.~Okafor:} methodology, software, investigation,
visualization. \textbf{L.~Fischer:} supervision, funding acquisition,
writing -- review \& editing.

\section*{Competing interests}
The authors declare no competing interests.

\bibliographystyle{unsrtnat}
\bibliography{references}

% --- Appendix --------------------------------------------------------------

\appendix

\section{Proof of Corollary~\ref{cor:radius}}\label{app:proof}

Let $i$ be the top-1 expert at $x$, so that $g_{\tau,i}(x) - g_{\tau,j}(x)
\ge m(x)$ for every $j \ne i$, and write $\Delta = g_\tau(y) - g_\tau(x)$.
For every $j \ne i$,
\begin{align*}
  g_{\tau,i}(y) - g_{\tau,j}(y)
    &= g_{\tau,i}(x) - g_{\tau,j}(x) + \Delta_i - \Delta_j \\
    &\ge m(x) - \lvert \Delta_i - \Delta_j \rvert .
\end{align*}
Since $(\Delta_i - \Delta_j)^2 \le 2(\Delta_i^2 + \Delta_j^2)
\le 2\norm{\Delta}_2^2$, Theorem~\ref{thm:lipschitz} gives
\[
  \lvert \Delta_i - \Delta_j \rvert \le \sqrt{2}\, \norm{\Delta}_2
  \le \frac{\norm{W}_2}{\sqrt{2}\, \tau}\, \norm{x - y}_2 .
\]
The right-hand side is smaller than $m(x)$ whenever
$\norm{x - y}_2 < \sqrt{2}\, \tau\, m(x) / \norm{W}_2$. Then
$g_{\tau,i}(y) > g_{\tau,j}(y)$ for every $j \ne i$, so $i$ is also the
top-1 expert at $y$. \qed

\section{Training details}\label{app:details}

Table~\ref{tab:hyper} lists the configuration shared by all models in
Section~\ref{sec:experiments}.

\begin{table}[htb]
  \caption{Training configuration shared by all models.}
  \label{tab:hyper}
  \centering\small
  \begin{tabular}{@{}ll@{}}
    \toprule
    Setting & Value \\
    \midrule
    Layers, width, heads & 12, 768, 12 \\
    Experts per MoE layer & 8, in every other block \\
    Routing & top-2, capacity factor 1.25 \\
    Optimiser & Adam, $\beta = (0.9, 0.95)$ \\
    Learning rate & $6 \times 10^{-4}$, cosine decay \\
    Batch size & $256 \times 1{,}024$ tokens \\
    Training steps & 60{,}000 \\
    Balance weight $\alpha$ & $10^{-2}$ \\
    Target constant $L_\star$ & 4 \\
    Minimum temperature $\tau_{\min}$ & 0.25 \\
    \bottomrule
  \end{tabular}
\end{table}

\end{document}