papers/zen-mixture-of-experts.tex

\documentclass[11pt,a4paper]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{amsmath,amsfonts,amssymb}
\usepackage{graphicx}
\usepackage{hyperref}
\usepackage{listings}
\usepackage{color}
\usepackage{booktabs}
\usepackage{float}
\usepackage{geometry}
\geometry{margin=1in}
\definecolor{zenblue}{RGB}{41,121,255}
\hypersetup{colorlinks=true,linkcolor=zenblue,urlcolor=zenblue,citecolor=zenblue}

\title{\textbf{Zen MoDE: Mixture of Distilled Experts Architecture}\\
\large Technical Report v2025.06}
\author{Antje Worring, Zach Kelling \\ Zen LM Research Team\\
\texttt{[email protected]}}
\date{June 2025}

\begin{document}
\maketitle

\begin{abstract}
We present Zen MoDE (Mixture of Distilled Experts), the core architectural innovation underlying the Zen model family. Zen MoDE addresses two fundamental challenges in large language model scaling: computational efficiency via sparse expert routing, and knowledge transfer via distilled expert initialization. We introduce semantic routing, where experts specialize by input domain and linguistic style rather than by position-hashing or random assignment. Combined with distilled expert initialization from smaller dense teacher models, Zen MoDE achieves performance matching a 400B-parameter dense model using only 47B active parameters (480B total). We present the router architecture, expert specialization analysis, load balancing algorithm, and performance benchmarks comparing Zen MoDE against dense equivalents.
\end{abstract}

\section{Introduction}

Scaling language models requires navigating a fundamental tension: more parameters improve capability, but increase inference cost proportionally. Mixture-of-Experts (MoE) architectures offer an elegant resolution: maintain a large parameter count while activating only a sparse subset per forward pass.

Conventional MoE implementations route tokens to experts via a learned linear router, trained to maximize routing utility. In practice, this leads to load imbalance (popular experts become overloaded while others are underutilized) and poor expert specialization (experts do not develop interpretable specializations).

Zen MoDE introduces three innovations that address these limitations:

\begin{enumerate}
  \item \textbf{Semantic routing}: Expert selection guided by semantic content embeddings, enabling interpretable domain and style specialization.
  \item \textbf{Distilled expert initialization}: Each expert is initialized from a smaller dense teacher model specialized on a semantic domain, providing warm-start knowledge.
  \item \textbf{Adaptive capacity factors}: Dynamic capacity allocation based on real-time expert load, preventing token dropping under skewed distributions.
\end{enumerate}

\section{Architecture}

\subsection{Overall Structure}

Zen MoDE follows the standard Transformer architecture with MoE layers replacing every-other dense feed-forward network (FFN) layer:

\begin{itemize}
  \item \textbf{Layer structure}: Alternating dense attention + {dense FFN, MoE FFN} pattern.
  \item \textbf{Expert count}: 128 experts per MoE layer.
  \item \textbf{Active experts}: Top-$K=8$ experts per token ($K=2$ for the 14B model).
  \item \textbf{Expert FFN size}: $4d_{\text{model}}$ hidden dimension (same as dense FFN equivalent).
  \item \textbf{Shared experts}: 4 experts always activated (``shared'' experts, active for every token).
\end{itemize}

The total parameter count $P_{\text{total}}$ and active parameters $P_{\text{active}}$ for an $L$-layer model with $N_E$ experts:

\begin{align}
P_{\text{total}} &= P_{\text{attn}} + P_{\text{shared}} + N_E \cdot P_{\text{expert}} \\
P_{\text{active}} &= P_{\text{attn}} + P_{\text{shared}} + K \cdot P_{\text{expert}}
\end{align}

For the 480B Zen MoDE: $P_{\text{total}} = 480B$, $P_{\text{active}} = 47B$ (9.8\% activation ratio).

\subsection{Model Scale Configurations}

\begin{table}[H]
\centering
\caption{Zen MoDE model configurations}
\label{tab:configs}
\begin{tabular}{lcccccc}
\toprule
Model & Total Params & Active Params & Experts & $K$ & Layers & $d_{\text{model}}$ \\
\midrule
Zen MoDE 14B & 14B & 2.4B & 64 & 2 & 28 & 2048 \\
Zen MoDE 72B & 72B & 8.1B & 128 & 4 & 64 & 4096 \\
Zen MoDE 236B & 236B & 22B & 128 & 8 & 80 & 7168 \\
Zen MoDE 480B & 480B & 47B & 128 & 8 & 94 & 7680 \\
\bottomrule
\end{tabular}
\end{table}

\section{Semantic Routing}

\subsection{Router Architecture}

Conventional MoE routers apply a linear projection to the token representation and take the top-$K$ softmax scores:

\begin{equation}
\text{router}_{\text{standard}}(x) = \text{TopK}(\text{softmax}(\mathbf{W}_r x), K)
\end{equation}

The Zen MoDE semantic router adds a two-stage process:

\paragraph{Stage 1: Semantic Embedding.} A lightweight embedding model $f_s$ maps the input sequence prefix to a 256-dimensional semantic vector:
\begin{equation}
\mathbf{s} = f_s(\text{prefix}) \in \mathbb{R}^{256}
\end{equation}

This semantic vector captures topic domain, linguistic register, and task type at the sequence level.

\paragraph{Stage 2: Expert Scoring.} Token-level routing is biased toward experts whose learned prototype $\mathbf{p}_e$ is nearest to the semantic vector:

\begin{equation}
\text{score}(x, e) = \underbrace{\mathbf{w}_e^\top x}_{\text{token routing}} + \underbrace{\beta \cdot \text{sim}(\mathbf{s}, \mathbf{p}_e)}_{\text{semantic bias}}
\end{equation}

where $\beta = 0.3$ balances local (token) and global (semantic) routing signals.

\subsection{Expert Specialization}

\begin{table}[H]
\centering
\caption{Expert specialization analysis: top domains per expert cluster (128 experts, 16 clusters)}
\label{tab:specialization}
\begin{tabular}{lll}
\toprule
Cluster & Primary Domain & Secondary Domains \\
\midrule
1--8 & Code (Python, JS, Rust) & Algorithms, data structures \\
9--16 & Mathematics (formal) & Proofs, theorems \\
17--24 & Scientific writing & Academic, citations \\
25--32 & Medical/clinical & Pharmacology, anatomy \\
33--40 & Legal text & Contracts, statutes \\
41--48 & Financial text & Earnings, risk, regulation \\
49--56 & Creative writing & Fiction, poetry \\
57--64 & Conversational & Dialogue, casual register \\
65--72 & News and journalism & Current events, AP style \\
73--80 & Technical documentation & Manuals, API docs \\
81--88 & Multilingual (CJK) & Japanese, Chinese, Korean \\
89--96 & Multilingual (Arabic) & Arabic, Farsi, Hebrew \\
97--104 & Multilingual (Slavic) & Russian, Polish, Czech \\
105--112 & Low-resource languages & African, SE Asian \\
113--120 & Instruction following & Commands, structured output \\
121--128 & Reasoning (chain-of-thought) & Logic, planning \\
\bottomrule
\end{tabular}
\end{table}

Expert specialization is measured by the Jensen-Shannon divergence between the token distribution of expert $e$'s input and the corpus-wide distribution. Higher divergence indicates stronger specialization. Zen MoDE experts achieve average JS-divergence of 0.42 versus 0.18 for standard top-K routing (2.3$\times$ improvement).

\section{Distilled Expert Initialization}

\subsection{Teacher Model Preparation}

For each semantic domain cluster $c$, we prepare a specialist dense teacher model:
\begin{enumerate}
  \item Start from a 7B dense model pre-trained on the full corpus.
  \item Fine-tune on domain-specific data (e.g., code corpus for code experts, PubMed for medical experts).
  \item Extract the FFN weight matrices $\{W_1^c, W_2^c\}$ for each layer as expert initialization.
\end{enumerate}

\subsection{Expert Initialization Protocol}

Expert $e$ in cluster $c$ is initialized by projecting the teacher FFN weights to expert dimensions:

\begin{equation}
W_{e,\text{init}} = \mathbf{P}_e \cdot W^c_{\text{teacher}} \cdot \mathbf{Q}_e
\end{equation}

where $\mathbf{P}_e, \mathbf{Q}_e$ are random orthogonal projections that map from teacher dimensions to expert dimensions. This preserves the spectral properties of the teacher weights in a lower-dimensional subspace.

\begin{table}[H]
\centering
\caption{Distilled initialization vs. random initialization (perplexity on held-out eval, lower is better)}
\label{tab:distillation}
\begin{tabular}{lccc}
\toprule
Initialization & Step 1K & Step 10K & Step 100K \\
\midrule
Random & 48.2 & 18.4 & 8.4 \\
Copy from dense base & 22.4 & 12.8 & 7.8 \\
Distilled (ours) & \textbf{16.4} & \textbf{10.2} & \textbf{7.2} \\
\bottomrule
\end{tabular}
\end{table}

Distilled initialization provides significant early-training advantage, converging to final loss in 31\% fewer training steps.

\section{Load Balancing}

\subsection{Expert Load Imbalance Problem}

Without load balancing, popular experts become overloaded while others are underutilized. Overloaded experts either drop tokens (capacity-limited) or become bottlenecks (unlimited capacity). Both outcomes degrade performance.

\subsection{Auxiliary Load Balancing Loss}

Zen MoDE minimizes an auxiliary load balancing loss alongside the main language modeling objective:

\begin{equation}
\mathcal{L}_{\text{balance}} = \alpha \cdot N \sum_{e=1}^{N_E} f_e \cdot P_e
\end{equation}

where $f_e = \frac{1}{T} \sum_t \mathbf{1}[e \in \text{TopK}(t)]$ is the fraction of tokens routed to expert $e$, $P_e = \frac{1}{T} \sum_t p_t^e$ is the average routing probability for expert $e$, and $\alpha = 0.01$ is the load balancing coefficient.

\subsection{Adaptive Capacity Factors}

The capacity factor $C$ limits the maximum tokens an expert can process per batch:
\begin{equation}
\text{capacity}_e = C \cdot \frac{\text{batch size} \cdot \text{seq len}}{N_E}
\end{equation}

Zen MoDE uses adaptive capacity factors that adjust per-expert based on rolling load statistics:
\begin{equation}
C_e(t) = C_{\text{base}} + \gamma \cdot \left(\bar{f}_e - \frac{1}{N_E}\right)
\end{equation}

Under-utilized experts ($\bar{f}_e < 1/N_E$) get reduced capacity, over-utilized experts get increased capacity, preventing token dropping under natural load distribution skew.

\section{Benchmark Results}

\subsection{Performance vs. Dense Equivalent}

\begin{table}[H]
\centering
\caption{Zen MoDE vs. dense equivalent: MMLU 5-shot (\%)}
\label{tab:perf_dense}
\begin{tabular}{lccc}
\toprule
Model & Params & Active & MMLU \\
\midrule
Dense 7B & 7B & 7B & 63.4 \\
Zen MoDE 14B & 14B & 2.4B & 72.8 \\
Dense 13B & 13B & 13B & 67.2 \\
Dense 70B & 70B & 70B & 80.2 \\
Zen MoDE 72B & 72B & 8.1B & 84.8 \\
Dense 400B (estimated) & 400B & 400B & 90.4 \\
Zen MoDE 480B & 480B & 47B & 91.2 \\
\bottomrule
\end{tabular}
\end{table}

Zen MoDE 480B matches estimated dense-400B performance with only 47B active parameters—an 8.5$\times$ reduction in inference FLOPs.

\subsection{Expert Utilization}

\begin{table}[H]
\centering
\caption{Expert utilization metrics (Zen MoDE 72B)}
\label{tab:utilization}
\begin{tabular}{lcc}
\toprule
Metric & Standard Top-K & Zen MoDE \\
\midrule
Utilization entropy & 5.42 bits & 6.81 bits (max 7 bits) \\
CV (coefficient of variation) & 0.84 & 0.31 \\
Token drop rate & 3.2\% & 0.8\% \\
Expert utilization below 50\% & 24 experts & 8 experts \\
Expert utilization above 200\% & 18 experts & 3 experts \\
\bottomrule
\end{tabular}
\end{table}

\subsection{Routing Quality}

\begin{table}[H]
\centering
\caption{Routing quality metrics}
\label{tab:routing}
\begin{tabular}{lcc}
\toprule
Metric & Standard Top-K & Semantic Routing \\
\midrule
Domain consistency (same-domain token routing) & 61.4\% & 84.2\% \\
Expert JS-divergence & 0.18 & 0.42 \\
Router confidence (avg top score) & 0.48 & 0.62 \\
Cross-domain routing (noise) & 38.6\% & 15.8\% \\
\bottomrule
\end{tabular}
\end{table}

\section{Training Efficiency}

\subsection{Compute Budget}

The MoDE architecture provides significant training compute savings over dense models at equivalent capacity:

\begin{equation}
\frac{\text{FLOPs}_{\text{MoDE}}}{\text{FLOPs}_{\text{dense-equivalent}}} = \frac{K + N_{\text{shared}}}{N_E} \approx \frac{12}{128} \approx 9.4\%
\end{equation}

In practice, attention layers (non-sparse) and overhead reduce the realized savings to approximately 8$\times$ reduction in training FLOPs for the same total parameters.

\subsection{Communication Overhead (Distributed Training)}

Expert parallelism requires all-to-all communication to route tokens to the correct expert device. Communication overhead as a fraction of total training time:

\begin{table}[H]
\centering
\caption{Expert parallelism communication overhead}
\label{tab:comm}
\begin{tabular}{lccc}
\toprule
Model & Expert Parallelism & Comm / Total & Effective Throughput \\
\midrule
72B & 8 nodes & 8.4\% & 91.6\% of compute \\
236B & 16 nodes & 11.2\% & 88.8\% of compute \\
480B & 32 nodes & 14.8\% & 85.2\% of compute \\
\bottomrule
\end{tabular}
\end{table}

\section{Conclusion}

Zen MoDE demonstrates that principled mixture-of-experts design—semantic routing, distilled expert initialization, and adaptive load balancing—substantially improves over naive sparse routing approaches. The 480B Zen MoDE model matches estimated dense-400B performance at 8.5$\times$ lower inference FLOPs, and semantic routing increases expert specialization by 2.3$\times$ as measured by domain-routing consistency. Distilled expert initialization accelerates convergence by 31\%, enabling efficient training of very large expert counts.

\begin{thebibliography}{99}
\bibitem{moe_original} Shazeer, N. et al. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. \textit{ICLR}, 2017.
\bibitem{switch} Fedus, W. et al. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. \textit{JMLR}, 2022.
\bibitem{gshard} Lepikhin, D. et al. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. \textit{ICLR}, 2021.
\bibitem{mixtral} Jiang, A. et al. Mixtral of Experts. \textit{arXiv:2401.04088}, 2024.
\bibitem{expert_specialization} Zuo, S. et al. Taming Sparsely Activated Transformer with Stochastic Experts. \textit{ICLR}, 2022.
\end{thebibliography}

\end{document}