JetMoE: Reaching Llama2 Performance with 0.1M Dollars

🎯 Pitch

A sparse mixture-of-experts model matching Llama2-7B and beating Llama2-13B-Chat was trained entirely on public data for under $100,000—making competitive LLM development radically more accessible. The key twist: extending sparse activation to both attention heads and feed-forward layers, activating just 2B of its 8B parameters per token.

1. Executive Summary

This paper introduces JetMoE-8B, a sparsely-gated mixture-of-experts language model that extends sparse activation to both attention and feed-forward layers—using Mixture of Attention heads (MoA) for query and output projections with shared key/value matrices, and a top-2 routed SMoE for the MLP—trained entirely on public datasets for under $0.1 million (30,000 H100 GPU hours, 1.25T tokens). JetMoE-8B achieves 8B total parameters while activating only 2B per token, outperforming the dense Llama2-7B on the OpenLLM Leaderboard (53.0 vs. 51.0 average) and, after distilled SFT and distilled DPO alignment, surpassing Llama2-13B-Chat on MT-Bench (6.48 vs. 6.27). The paper establishes that competitive LLM performance can be reached at dramatically lower cost, though the single budget-constrained training run precludes ablation of the architecture’s individual contributions—leaving open which design choices (MoA, shared KV, two-phase data curriculum) are responsible for the gains.

2. Context and Motivation

The Core Problem: The Resource Barrier to Competitive LLM Development

The fundamental problem this paper addresses is the prohibitively high cost of training large language models, which creates a structural barrier that excludes academic researchers, small companies, and independent developers from participating in frontier LLM research. While the executive summary notes that JetMoE-8B was trained for under $0.1 million, this figure acquires meaning only when situated against the broader economic context the paper implies: training a model like Llama2-7B or its equivalents typically requires millions of dollars in compute and access to massive proprietary datasets that only large corporations can assemble.

This resource concentration is not merely an inconvenience—it shapes what research questions get asked, what architectures get explored, and who gets to build the foundation models that increasingly underlie downstream applications. The paper frames this implicitly through its emphasis on openness:

"JetMoE-8B is highly open and academia-friendly, using only public datasets and training code. All training parameters and data mixtures have been detailed in this report to facilitate future efforts in the development of open foundation models."

The word "facilitate" is doing real work here. The authors are not just releasing a model—they are attempting to lower the barrier to entry so that subsequent work can build on their recipe without needing to replicate the costly exploration phase.

The Specific Gap: Efficient Architectures Without the Efficiency of Full Transparency

The paper identifies a gap between two trends in the MoE literature. On one side, major commercial MoE models (Deepseek MoE, Mixtral 8x7B, Grok-1, DBRX) have demonstrated that sparse architectures can achieve excellent performance with reduced inference compute. However, as the paper notes:

"even though these models achieve excellent performance, they are not truly open-sourced as the training recipes are not published and may contain proprietary datasets inaccessible outside of large corporations."

This is a specific, concrete complaint. The models themselves may be released as weights, but the recipe—the data mixture, the training infrastructure decisions, the hyperparameter selection process, the debugging of training instabilities—remains proprietary. For a researcher wanting to understand why a particular MoE design choice works, or wanting to experiment with modifications, having weights is insufficient. You need to be able to reproduce the training process to run ablations, vary components, and understand causality. Without the recipe, these commercial models are essentially un-reproducible black boxes from a research perspective.

On the other side, open-source MoE attempts have struggled to demonstrate competitive performance. The paper cites OpenMoE specifically:

"The open-source community has also attempted to train MoE models, such as OpenMoE... but its performance is only on par with weak dense models with similar activation parameters, such as OpenLLaMA and TinyLLaMA."

This is a damning comparison. If an MoE model's performance merely matches a dense model with the same number of active parameters, then the entire point of the MoE architecture—scaling total parameters while keeping inference cost constant—is undermined. The MoE should be better than a comparably-sized dense model because it has access to more total parameters through the expert routing. The fact that open-source MoE attempts have failed to demonstrate this advantage suggests that making MoEs work well requires non-obvious design choices that the commercial labs have discovered but not disclosed.

This creates a specific, actionable gap: there exists no publicly documented, reproducible training run that demonstrates how to build a competitive MoE model from scratch using only open resources. JetMoE-8B is positioned to fill exactly this gap.

The Architectural Challenge: Why Don't MoEs "Just Work"?

To understand the paper's motivation, it's worth unpacking why successful MoE training is hard, which the paper addresses through its design choices rather than through explicit discussion. MoEs introduce several failure modes that don't exist in dense models:

Load imbalance. The router might consistently favor a few experts, leaving the rest underutilized. This is not just an efficiency concern—it means the model is essentially not using the extra capacity it was given, collapsing back toward dense behavior. The paper cites the standard solution (frequency-based auxiliary loss from Fedus et al., 2021) but also introduces router z-loss from Zoph et al. (2022), suggesting that standard load balancing alone was insufficient for training stability.

Training instability. The router's discrete selection (top-k) creates a non-differentiable operation in the computation graph. Gradients only flow to the selected experts, and if routing decisions fluctuate during training, the effective optimization landscape can be noisy. The paper's use of z-loss—which penalizes large logit magnitudes in the router—is specifically motivated by stabilizing training, since large router logits create near-deterministic expert assignments that prevent gradient signal from reaching underutilized experts.

The attention bottleneck. In most MoE architectures (Deepseek MoE, Mixtral, OpenMoE), only the feed-forward layers are replaced with MoE layers. The self-attention layers remain dense. This means that as you scale up the FFN capacity with more experts, the attention layers become a bottleneck—they consume a fixed fraction of the total compute but don't benefit from the parameter scaling. The paper's key architectural contribution—extending sparsity to attention layers via Mixture of Attention heads (MoA)—addresses this directly:

"By introducing the MoA, we can scale up the attention layer with more attention experts while maintaining the same amount of computation. Such that the attention layer will not become a performance bottleneck, while we scale up the MLP layers."

This is a subtle but important motivation. If you only sparsify the FFN, then at some scale the dense attention layers dominate the compute budget and limit further efficiency gains. The MoA design allows both attention and FFN to scale together, maintaining a balanced architecture.

The Data Transparency Gap

Beyond architecture, the paper identifies a second transparency gap: training data composition. Commercial LLMs are typically trained on proprietary data mixtures that are carefully curated but never disclosed. This makes it impossible to determine whether a model's performance comes from architectural innovation or from having access to better training data—a classic confound in ML research.

The paper addresses this by being exhaustively specific about its data sources (Sections 3.1 and 3.2, tabulated in Table 2) and its two-phase training curriculum (Figure 2). The inclusion of synthetic datasets (OpenHermes 2.5, UltraTextbooks, Magicoder) alongside traditional web crawl data (RefinedWeb, StarCoder, The Pile) is a deliberate choice that reflects an emerging consensus—visible in miniCPM and Gemma—that data quality and diversity matter more than raw quantity, especially toward the end of training. By documenting the exact mixture, the paper enables other researchers to ablate data choices and understand their impact, something impossible with proprietary models.

How the Paper Positions Itself

The paper positions JetMoE-8B as a demonstration of possibility rather than a claim of architectural superiority. It explicitly disclaims any ablation studies:

"Due to the limited $100k budget, we can not afford any ablation study for the model architecture. The hyperparameters and data mixtures are also handpicked based on the empirical results from previous works"

This is an unusual and honest admission. Most papers present their design choices as justified by evidence; this paper presents them as conjectures that happen to work well together but whose individual contributions are unknown. The positioning is: "Here is a recipe that produces a competitive model at a fraction of the expected cost. We don't know exactly why every ingredient matters, but we document everything so others can find out."

This framing serves a specific strategic purpose in the research ecosystem. The paper is not claiming to have solved MoE training—it is claiming to have demonstrated that the barrier to entry is lower than commonly assumed, and it is providing a fully-reproducible starting point that lowers that barrier for everyone else. The $0.1 million figure is not just a cost report; it is a challenge to the field's implicit assumption that competitive LLMs require massive corporate resources.

The comparison targets in Table 3 are carefully chosen to support this narrative: JetMoE-8B doesn't claim to beat the best models overall. It claims to beat Llama2-7B, a widely-used open model that represents a reasonable baseline for what the community considers a "good" 7B dense model. The fact that JetMoE-8B achieves this with 1.25T tokens (vs. Llama2's 2T), fewer active parameters (2B vs. 7B), and lower cost reinforces the efficiency argument without making unrealistic absolute performance claims.

Similarly, the MT-Bench comparison in Table 4 shows JetMoE-8B-Chat (6.48) surpassing Llama2-13B-Chat (6.27) but still trailing GPT-3.5-turbo (8.14). This positions the model as "competitive with the best open chat models at its size class" rather than "state-of-the-art overall"—an honest and strategically appropriate claim.

3. Technical Approach

3.1 Reader Orientation

This is primarily an engineering contribution paper that describes the design and training of JetMoE-8B, a large language model, with the core idea being that competitive LLM performance can be achieved at dramatically lower cost by applying sparse activation to both attention and feed-forward layers within a Mixture-of-Experts framework, trained on a carefully curated, fully open-source data mixture using a two-phase curriculum. The system is a sparsely-gated Transformer where every transformer block contains an MoE for the MLP and an MoE for the attention mechanism, allowing the model to have 8 billion total parameters but activate only 2.2 billion for any given input token into the selected experts.

The problem it solves is the prohibitively high computational cost and closed nature of training frontier LLMs, and the solution is shaped as a fully open-source recipe—model architecture, training data, infrastructure choices, and hyperparameters are all disclosed—that demonstrates a specific path to building a capable model for under $0.1 million, thereby lowering the barrier to entry for academic and independent researchers.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components, described in the order they are encountered when training and deploying the model:

The JetMoE Transformer Body: A stack of Transformer decoder layers where the standard dense feed-forward network (FFN) is replaced by a Sparsely-gated Mixture-of-Experts (SMoE) layer, and the dense self-attention is replaced by a Mixture of Attention heads (MoA) layer. This is the core architectural innovation.
The Sparse Router: A lightweight, learned linear layer present in both the SMoE and MoA layers. Its job is to look at each input token representation and, using a top- $k$ selection mechanism, dynamically choose a small subset of experts to process that token, ensuring the total compute per token stays low.
The Experts: A collection of identical, independent sub-networks. For the SMoE, these are standard SwiGLU MLP layers. For the MoA, these are independent multi-head attention modules that share key and value projections but have their own query and output projections.
The Load Balancing and Stability Mechanisms: Auxiliary loss functions (a frequency-based load balancing loss and a router z-loss) that are added to the main language modeling loss. Their purpose is to encourage the router to use all experts evenly and to prevent the router's internal values from exploding, which would cause training instability.
The Two-Phase Training Curriculum: A data and learning rate schedule. Phase 1 uses a broad web-scale data mixture with a constant high learning rate. Phase 2 continues training with a decaying learning rate and a modified, higher-quality data mixture to refine the model's knowledge.

Information flows as follows: a sequence of text tokens is embedded into vectors → the sequence passes through a series of JetMoE layers. In each layer, every token's vector first goes to the MoA router, which selects 2 of 8 attention experts. The token is processed by these experts' attention heads to produce a contextualized representation. This output then goes to the SMoE router, which selects 2 of 8 MLP experts to further process it. The outputs from the selected experts are combined via a weighted sum, and the result is passed to the next layer. The final layer's output is used to predict the next token in the sequence.

3.3 Roadmap for the Deep Dive

First, the foundational SMoE mechanism: understanding the routing framework (§3.4.1) is essential because everything else is a specific instance of it. We'll define the router, the top- $k$ selection, and the weighted combination.
Second, the Feed-Forward Expert (SMoE): this is the standard application of MoE, replacing the dense MLP in a transformer block with a sparsely activated one (§3.4.2).
Third, the Attention Expert (MoA): this is a key architectural innovation. We'll dissect how the attention mechanism is decomposed into experts and the crucial "shared $k, v$ " design choice that aims to improve training stability and efficiency (§3.4.3).
Fourth, the load balancing and stability losses: we'll formalize the auxiliary losses that make the sparse training work in practice, preventing expert collapse and numerical instability (§3.4.4).
Fifth, the training recipe and two-phase data curriculum: this section translates the architecture into a concrete training run, covering the infrastructure, standard hyperparameters, learning rate schedule, and the precise data mixtures used in each phase (§3.4.5).

3.4 Detailed, Sentence-Based Technical Breakdown

Sparsely-Gated Mixture-of-Experts (SMoE) Routing Framework

The core mechanism that enables JetMoE's efficiency is the sparsely-gated Mixture-of-Experts layer, which replaces a standard dense layer. Instead of sending every input token through a single, large, computationally expensive function, an SMoE layer contains a set of $N$ smaller, independent modules called experts, denoted $f_1, \ldots, f_N$ , and a router that decides which specific experts should process each token.

The router computes a probability distribution over the experts, and a top-$k$ selection mechanism activates only the $k$ experts with the highest probability, where $k < N$ . This means that even though the layer contains $N$ experts' worth of total parameters, only $k$ of them are used for any single input token, making the per-token computational cost similar to a single, small expert module.

The router in JetMoE is implemented as a simple linear layer. Given an input token representation $\mathbf{x} \in \mathbb{R}^{D_{\text{emb}}}$ , a matrix of expert embeddings $\mathbf{W}_{rtr} \in \mathbb{R}^{N \times D_{\text{emb}}}$ is used to calculate a raw score for each expert.

$\mathbf{s} = \mathbf{W}_{rtr}\mathbf{x}$

where $\mathbf{s} \in \mathbb{R}^N$ is a vector of router logits, one score for each of the $N$ experts, $\mathbf{W}_{rtr}$ is a learned weight matrix that projects the input to a score for each expert, and $\mathbf{x}$ is the input token's hidden representation.

What it computes: a simple dot-product affinity between the input token's representation and a learned "prototype" vector for each expert. The higher the score $s_i$ , the more the router believes expert $i$ is suited for this token. The output is a vector of unnormalized scores.

Why this form: a linear projection is chosen for its speed and simplicity, as the routing decision must be made for every token in the sequence. A more complex router (e.g., an MLP) would consume significant computation, potentially negating the efficiency savings of having multiple lightweight experts.

Next, the raw scores $\mathbf{s}$ are converted into a sparse probability vector $g(e \mid \mathbf{x})$ that assigns a non-zero weight only to the top- $k$ experts. This is done by first applying the $\text{Top}k$ operator to zero out all but the top- $k$ logits, and then applying a softmax function over the surviving logits.

g(e \mid \mathbf{x}) = \begin{cases} \mathrm{softmax}\left(\text{Top}k\left(\mathbf{s}\right)\right)_i, & \mathbf{s}_i \in \text{Top}k\left(\mathbf{s}\right) \\ 0, & \mathbf{s}_i \notin \text{Top}k\left(\mathbf{s}\right) \end{cases}

where $\text{Top}k(\mathbf{s})$ is an operation that retains only the largest $k$ values in the vector $\mathbf{s}$ , setting all other elements to $-\infty$ before the softmax, $\mathrm{softmax}(\cdot)_i$ is the standard softmax function applied to the sparsified logits, and $g(e \mid \mathbf{x})$ is the final routing weight assigned to expert $e$ for input $\mathbf{x}$ . If expert $e$ is not in the top- $k$ , its routing weight is flatly zero.

What it computes: a sparse, non-negative vector of length $N$ that sums to 1, with exactly $k$ non-zero entries. It acts as a hard gating mechanism, assigning a portion of the input's "attention" to the most relevant experts and completely ignoring the rest. The softmax ensures the selected experts' weights are in a normalized probability space.

Why this form: Zeroing out the non-top- $k$ experts is essential for the computational savings. If the weights were small but non-zero, we would still have to compute $f_e(\mathbf{x})$ for all experts, which defeats the purpose of a sparse MoE. The softmax on the top- $k$ logits provides a differentiable way to assign relative importance among the chosen experts, allowing gradients to flow back to the router for training.

Finally, the output of the entire SMoE layer is the weighted sum of the outputs from the selected experts.

$y = \sum_{e=1}^{N} g(e \mid \mathbf{x}) \cdot f_e(\mathbf{x})$

where $y$ is the layer's output for the input token $\mathbf{x}$ , and $f_e(\mathbf{x})$ is the output of the $e$ -th expert module.

What it computes: a combination of the outputs from the $k$ chosen experts, weighted by the router's confidence in each expert for that token. Because $g(e \mid \mathbf{x})$ is zero for $N-k$ experts, the sum is effectively only over $k$ terms, and the unselected experts' functions $f_e(\mathbf{x})$ are never evaluated.

Why this form: A weighted sum is a simple, differentiable aggregation method. It allows the router to not only choose which experts to use but also modulate their contribution, creating a richer combinatorial space for the output representation. For instance, a token could be 70% processed by a "general language" expert and 30% by a "mathematical reasoning" expert.

Feed-Forward Expert (SMoE)

In JetMoE, each standard feed-forward network (FFN) transformer block is replaced by an SMoE layer with $N=8$ experts and a top- $k=2$ gating mechanism. Each individual expert in this set is a standard SwiGLU-activated, 2-layer MLP with a hidden size $D_{\mathrm{ffd}}$ .

$f_{mlp}(\mathbf{x}) = \mathbf{W}_{out} \cdot \sigma\left(\mathbf{W}_{in} \mathbf{x}\right)$

where $\mathbf{W}_{in}$ is the input projection matrix of shape $(2D_{\mathrm{ffd}}, D_{\mathrm{emb}})$ , $\mathbf{W}_{out}$ is the output projection matrix of shape $(D_{\mathrm{emb}}, D_{\mathrm{ffd}})$ , and $\sigma$ is the SwiGLU activation function.

What it computes: The input token representation $\mathbf{x}$ is first projected to a higher dimension $2D_{\mathrm{ffd}}$ by $\mathbf{W}_{in}$ . The SwiGLU activation $\sigma$ applies a gating mechanism on this high-dimensional vector, effectively multiplying half of the dimensions by a sigmoid of the other half. The result is then projected back down to the model's embedding dimension $D_{\mathrm{emb}}$ by $\mathbf{W}_{out}$ . This is the standard MLP operation, and in the SMoE context, this operation is computed only by the 2 experts selected by the router for a given token.

Why this form: SwiGLU activation is a strong, empirically validated choice for large language models that consistently outperforms standard ReLU or GeLU activations. Replacing a single large FFN with 8 smaller SwiGLU FFNs (each with fewer hidden dimensions) and sparsely activating 2 of them creates a model with vastly more total FFN capacity (8 experts) for the same per-token inference cost as a dense model with similar dimensions to a single expert. The paper sets top- $k=2$ to give the model two degrees of freedom to combine different kinds of non-linear processing for each token, a common design choice for balancing performance and efficiency.

Attention Expert (Mixture of Attention Heads - MoA)

This is the paper's most significant architectural deviation from standard MoE transformers. Instead of only replacing the FFN with an MoE, JetMoE also applies the SMoE framework to the multi-head self-attention (MHA) module. This creates a set of $N$ attention experts, allowing the model to scale the capacity of the attention layers alongside the FFN layers without increasing per-token cost.

Each attention expert $e$ operates as an independent multi-head attention module. It has its own query projection matrix $\mathbf{W}_{q}^{e}$ and its own output projection matrix $\mathbf{W}_{o}^{e}$ . Critically, the key projection matrix $\mathbf{W}_{k}$ and the value projection matrix $\mathbf{W}_{v}$ are shared across all attention experts in the same layer.

For a given input sequence $\mathbf{x}$ , the shared $k$ and $v$ projections are computed once.

\mathbf{k} = \mathbf{W}_{k}\mathbf{x}

\mathbf{v} = \mathbf{W}_{v}\mathbf{x}

where $\mathbf{k}$ is the key vector sequence shared by all attention experts in this layer, and $\mathbf{v}$ is the value vector sequence shared by all attention experts.

What it computes: This step pre-computes the keys and values that will be accessed by every attention expert. It's the standard first step of multi-head attention, but done globally for the layer rather than inside each expert. The output is a single set of key and value tensors.

Why this form: Sharing key and value projections is a design choice explicitly made for training stability and inference efficiency. Computing separate keys and values for 8 attention experts would require 8 separate attention matrix multiplications between the queries and keys, and 8 separate weighted sums of values, dramatically increasing the memory bandwidth pressure and the computational cost of the attention module. Sharing $k$ and $v$ turns the $N$ -expert MoA block's cost into something much closer to a single, slightly wider attention block, making the sparse activation strategy practically feasible.

Inside a specific attention expert $e$ , its unique query $\mathbf{q}_e$ is computed, and then standard multi-head attention is performed using its unique query but the globally shared keys and values.

\mathbf{q}_{e} = \mathbf{W}_{q}^{e}\mathbf{x}

\mathbf{a}_{e} = \mathrm{MHA}\left(\mathbf{q}_{e}, \mathbf{k}, \mathbf{v}\right)

\mathbf{o}_{e} = \mathbf{W}_{o}^{e}\mathbf{a}_{e}

where $\mathbf{q}_e$ is the query vector sequence produced by expert $e$ 's unique query projection, $\mathbf{a}_e$ is the output of the multi-head attention operation, which computes attention scores between $\mathbf{q}_e$ and the shared $\mathbf{k}$ and uses them to perform a weighted sum over the shared $\mathbf{v}$ , $\mathbf{o}_e$ is the final output of the attention expert, projected back to the model dimension by expert $e$ 's unique output projection, and RoPE relative position embeddings are applied during the attention calculation.

What it computes: Even though all experts look at the same key and value information (the "memory" of the sequence), each expert asks a different "question" (query) of that memory. One attention expert might learn to focus on the immediate preceding tokens, another on long-range dependencies, and a third on specific syntactic structures. The output is a token's contextualized representation according to that expert's specialized attention pattern.

Why this form: This design creates a powerful form of attention specialisation without increasing the memory bandwidth cost of the key-value cache during inference. In a standard MoE transformer, the attention layers are a dense computational bottleneck. By making attention sparse, the paper claims that "the attention layer will not become a performance bottleneck, while we scale up the MLP layers," meaning the model's architecture can be balanced as it scales. Each attention expert in JetMoE contains $H$ attention heads, making $D_{\text{att}} = H \times D_{\text{head}}$ , and the total number of parameters in attention layers can be increased by adding more experts without increasing the per-token activation count beyond the selected $k$ experts.

The MoA router uses the same top-$k$ gating mechanism described in Section 3.4.1. Only the selected $k$ experts' queries $\mathbf{q}_e$ and output projections $\mathbf{W}_o^e$ are computed, while the shared key and value projections are always computed once.

Load Balancing and Stability Mechanisms

Sparse MoEs are notoriously difficult to train because the router can degenerate. Two primary failure modes are addressed by JetMoE with a pair of auxiliary loss functions added to the primary language modeling loss $loss_{lm}$ .

Failure Mode 1: Expert Collapse. The router could always select the same small set of highly-proficient experts for all tokens, completely ignoring the others. This wastes the model's capacity and effectively turns it back into a dense model with an unnecessary overhead. To combat this, JetMoE uses a frequency-based load balancing auxiliary loss called $loss_b$ .

$loss_b = N \sum_{i=1}^{N} f_i P_i$

where $N$ is the total number of experts, $f_i$ is the fraction of tokens in the current batch that are dispatched to expert $i$ , and $P_i$ is the fraction of the total router probability mass that is allocated to expert $i$ .

What it computes: A scalar penalty. If the selections are perfectly balanced, meaning each expert receives approximately $1/N$ fraction of the tokens ( $f_i \approx 1/N$ ) and has $1/N$ of the router probability mass ( $P_i \approx 1/N$ ), the loss is minimized. If a few experts get all the tokens ( $f_i$ is high) and all the probability ( $P_i$ is high), the product $f_i P_i$ becomes large, and the loss adds a high penalty to the training objective.

Why this form: This specific formulation is taken directly from the Switch Transformer (Fedus et al., 2021) paper and has become a standard, effective method. It encourages both uniform routing probability (from $P_i$ ) and uniform actual usage (from $f_i$ ). This dual pressure is important: a router could distribute probability evenly but still dispatch to a few experts if it is poorly calibrated; the $f_i$ term in the loss directly penalizes this behavior.

Failure Mode 2: Router Logit Explosion. The router logits $\mathbf{s}$ can grow in magnitude during training, leading to near-deterministic expert selection that is not only bad for load balancing but also creates numerical instability in the softmax. To combat this, JetMoE uses the router z-loss called $loss_z$ .

$loss_z = \frac{1}{B} \sum_{i=1}^{B} \left(\log \sum_{j=1}^{N} \exp(x_j^i)\right)^2$

where $B$ is the total number of tokens in the batch, $x_j^i$ is the raw router logit for the $j$ -th expert for the $i$ -th token, and $\exp(x_j^i)$ is the exponentiated logit.

What it computes: The inner sum $\sum_{j=1}^{N} \exp(x_j^i)$ is the log-sum-exp logit of the normalization constant for the softmax. The z-loss is the mean squared value of this log normalization constant across all tokens in the batch. It penalizes the router for producing large, unnormalized logits which would force the softmax to become a one-hot vector.

Why this form: The log-sum-exp function has a smoothing effect on its inputs. Penalizing its squared magnitude directly pushes the router to keep its internal logits on a numerically stable scale, avoiding the extremes where training can be both unstable and overly deterministic. This is a simple, scalar, and easily computed penalty, and it does not require any memory-intensive auxiliary state.

The final training loss is a weighted sum of these three components:

$loss = loss_{lm} + \alpha loss_b + \beta loss_z$

where $\alpha$ is the weight for the load balancing loss, set to 0.01, and $\beta$ is the weight for the z-loss, set to 0.001.

What it computes: A composite objective where the primary language model's perplexity is jointly optimized with the health of the sparse architecture. The hyperparameter values of 0.01 and 0.001 are taken from the ST-MoE (Zoph et al., 2022) recipe.

Why this form: A weighted sum is the simplest way to combine multiple objectives. The small weights reflect the fact that the language modeling loss $loss_{lm}$ is the primary goal; the auxiliary losses are a gentle but persistent regularization that guides the router away from pathological behaviors without dominating the optimization process.

The Training Recipe and Two-Phase Data Curriculum

The training process is a fixed recipe of infrastructure, hyperparameters, and a two-phase data and learning rate schedule, all chosen to get the model to converge within the 30,000 GPU-hour budget.

Infrastructure and Parallelism. Training is conducted on a cluster with 12 nodes, each containing 8 H100 GPUs (96 total GPUs), using the Megatron-LM framework integrated with MegaBlocks for efficient MoE computation. A crucial design choice is the use of pipeline parallelism instead of the more common expert parallelism for model distribution. The paper argues this is better because the Sparse MoE architecture has a narrower hidden state (reducing communication cost for pipeline parallelism) and because the dropless MoE schema can cause load imbalance across experts, which would make expert parallelism inefficient as some devices would be idle. Pipeline parallelism sidesteps this by computing all experts of a layer on a single device.

Standard Hyperparameters. The model architecture is based on a 1B-scale Transformer. The model is trained using the AdamW optimizer with $\beta_1=0.9$ , $\beta_2=0.95$ , a maximum learning rate of $5 \times 10^{-4}$ , a batch size of 4 million tokens with a sequence length of 4096, a weight decay of 0.1, and gradient clipping of 1.0. The model uses $N=8$ experts for every MoE layer and a top- $k=2$ routing strategy, effectively making it about twice the computational cost of a 1B dense model.

Warmup-Stable-Decay (WSD) Learning Rate Schedule. JetMoE employs a three-stage learning rate schedule that deviates from the standard cosine decay.

\begin{cases} \frac{s}{W} \ast \eta, & s < W \\ \eta, & W < s < S \\ f(s - S) \ast \eta, & S < s < S + D \end{cases}

where $s$ is the current training step, $\eta$ is the maximum learning rate ( $5 \times 10^{-4}$ ), $W$ is the number of steps at the end of a linear warmup phase (spanning 10 billion tokens), $S$ is the end of a stable, constant learning rate phase, $D$ is the number of steps in a final decay phase (spanning the last 250 billion tokens), and $f(s-S)$ is a decreasing function of $s$ that continues to reduce the learning rate to 10% of the maximum.

What it computes: A training trajectory that preserves a high, constant learning rate for most of the training process, potentially allowing more aggressive exploration of the loss landscape, before a final decay phase fine-tunes the model parameters.

Why this form: This WSD schedule is specifically referenced from the miniCPM paper. The rationale is that a stable, high learning rate phase can lead to better generalization than a schedule that decays too early. The final, relatively sharp annealing phase is intended to settle the model into a high-quality local minimum using high-quality data.

The Two-Phase Data Curriculum. The data strategy is divided into two distinct phases, directly correlated with the learning rate schedule. The total training data volume is 1.25 trillion tokens from a mixture of open-source, carefully documented datasets.

In Phase 1 (which encompasses the warmup and stable learning rate stages), the model is trained on a broad web-scale data mixture to build fundamental language understanding. The datasets are primarily RefinedWeb, Starcoder, The Pile, peS2o from Dolma, and OpenWebMath. The goal is to inject a diverse, massive amount of knowledge into the model.

In Phase 2 (which aligns with the learning rate decay stage), the data mixture is augmented with a much larger proportion of high-quality, curated datasets. This follows the approach in miniCPM and Gemma. The specific data mixture for Phase 2 is meticulously listed in Table 2 and includes:

Synthetic SFT data: UltraChat, OpenHermes, xP3x, Magicoder, Evol-Code Alpaca, and others.
Textbooks: UltraTextbooks.
High-quality math and code: TemplateGSM, StackMathQA, Code-290k-ShareGPT, and more.

What it computes: A structured learning process where the model first learns general knowledge from a broad but noisy signal, and then uses a decaying learning rate to "fine-tune" on a cleaner, more instructive signal in the final training steps. The exact percentages for each dataset in Phase 2 are provided in the paper's Table 2 (e.g., RefinedWeb at 39.8%, Starcoder at 19.6%, UltraTextbooks at 4.8%).

Why this form: This two-phase approach is an empirical finding from recent small-model training strategies. The intuition is that the high learning rate in Phase 1 prevents the model from overfitting to any one data source, and the stable high-rate period gives the model time to thoroughly process the large-scale general corpora. The decaying Phase 2 then teaches the model to prioritize and refine its knowledge based on high-quality patterns and instructions, essentially performing a form of curriculum learning where the quality of the "teacher signal" increases at the end of training.

4. Key Insights and Innovations

Innovation 1: Demonstrating That Competitive LLM Training Costs Can Plausibly Be Cut by an Order of Magnitude

The paper's most consequential contribution is not any single architectural trick but the existence proof that a model competitive with Llama2-7B can be trained for under $0.1 million—roughly an order of magnitude less than what the field implicitly assumed was necessary for this tier of performance. This is a contribution to the economics of LLM research rather than to model architecture, and it matters precisely because the cost barrier was so widely accepted that academic groups largely stopped attempting to train frontier models from scratch.

Prior to JetMoE, the dominant assumption—reinforced by the compute requirements reported for models like Llama2, Pythia, and even the open-source OpenMoE—was that training a competitive ~7B-scale model required millions of dollars in compute plus access to proprietary, carefully curated datasets. The paper challenges this assumption not through argument but through demonstration: it provides a fully documented recipe (architecture, data mixture, hyperparameters, infrastructure decisions) that produces a model beating Llama2-7B on the OpenLLM Leaderboard (53.0 vs. 51.0 average across six benchmarks) and a chat variant surpassing Llama2-13B-Chat on MT-Bench (6.48 vs. 6.27). The fact that this model activates only 2.2B parameters per token—about 70% fewer than Llama2-7B's 7B—makes the efficiency claim even starker: it's simultaneously cheaper, smaller at inference time, and more capable.

The significance of this finding extends beyond the specific model. It reframes the research landscape by demonstrating that the barrier to entry is not fundamental—it is a consequence of unexplored design space and insufficiently optimized training recipes. The paper explicitly acknowledges that its own recipe is likely suboptimal ("given the limited computing budget available, our data mixture might not be ideal") and that it performed no ablation studies. This frankness is strategically important: it invites the research community to treat the recipe as a starting point for optimization rather than as a claimed optimum. A ~10× cost reduction from an unoptimized recipe implies that further reductions may be possible once the community begins systematically studying the design space.

This is a fundamental reframing of what is possible, not an incremental refinement. The $0.1 million figure, combined with the fully open documentation, functions as a challenge to the field's implicit assumption that competitive LLMs require resources only available to large corporations.

However, the claim requires careful scoping. The paper compares against Llama2-7B (trained on 2T tokens, released in 2023) and Gemma-2B, not against more recent models like Llama3 or Mistral-7B-v3. The 1.25T token budget is well below the compute-optimal token count for an ~8B model. The result demonstrates a lower bound on what is achievable at this budget, not an upper bound on what the architecture could do with more resources. The paper's contribution is thus about democratizing access to the competitive range, not about pushing the frontier forward—a different, and arguably equally important, form of progress.

Innovation 2: Extending Mixture-of-Experts Sparsity to Attention Layers via MoA with Shared Key/Value Projections

The paper's core architectural move—applying the SMoE framework to both attention and feed-forward layers—is a genuine departure from the dominant MoE design pattern. Most MoE transformers (Switch Transformer, GLaM, Deepseek MoE, Mixtral, OpenMoE) replace only the FFN layers with MoE modules, keeping self-attention dense. The implicit assumption in these designs is that attention is already "efficient enough" or that making it sparse would introduce unacceptable instability or complexity.

JetMoE challenges this assumption through its Mixture of Attention heads (MoA) design, which makes each attention expert responsible for its own query and output projections while sharing key and value projections across all experts in a layer. This is a conceptual shift in how to think about attention specialisation in MoE models. The standard MoE approach treats experts as repositories of knowledge (via the FFN), with attention serving as a shared, general-purpose computation for contextualization. JetMoE's design treats attention as another dimension of specialisation: different experts can learn different attention patterns—one focusing on local syntax, another on long-range semantic dependencies, a third on specific entity tracking—while still sharing the underlying key-value memory.

The shared KV projection is the critical enabling insight. Without it, each attention expert would need to compute its own keys and values, multiplying the memory bandwidth pressure of the attention mechanism by the number of selected experts. The standard $N$ -expert MoA would require $k$ separate attention matrix multiplications and $k$ separate value-weighted sums, making the computational cost scale with $k$ rather than approaching constant cost. Sharing $k$ and $v$ decouples the attention capacity scaling from the memory bandwidth scaling: the query computation (cheap, per-expert) is what scales with the number of experts, while the key-value computation (expensive, per-token) remains constant regardless of how many experts exist.

This is a fundamental architectural innovation rather than an incremental tweak. It establishes a new axis of model scaling—attention capacity via expert count—that was previously unavailable in MoE designs. The paper's claim that "the attention layer will not become a performance bottleneck, while we scale up the MLP layers" captures a genuine design principle: in a model where FFN capacity grows with the number of experts, failing to scale attention proportionally creates an architectural imbalance that limits returns from additional FFN experts. MoA solves this by allowing both components to scale together.

The evidence for this innovation's effectiveness is necessarily indirect because the paper performs no ablation studies comparing MoA against a standard dense-attention MoE baseline. The model's overall performance (Table 3) demonstrates that the combined architecture works, but cannot disentangle whether the MoA component is responsible for the gains or whether the same performance could have been achieved with dense attention and more FFN experts. The paper acknowledges this limitation explicitly ("we can not afford any ablation study for the model architecture"), which means the MoA design claim should be understood as plausible and well-motivated rather than empirically validated by this paper. The innovation is in the design concept and its justification, not in a demonstrated ablation advantage.

Innovation 3: The Two-Phase Training Curriculum as a Practical Cost-Efficiency Lever for Small-Budget Training

The paper's data strategy implements a two-phase curriculum where Phase 1 uses broad web-scale data (RefinedWeb, Starcoder, The Pile) with a constant high learning rate, and Phase 2 switches to a higher-quality mixture (including synthetic instruction data, textbooks, and curated math/code datasets) while decaying the learning rate following the Warmup-Stable-Decay (WSD) schedule. This is not a novel idea in isolation—the paper explicitly credits miniCPM and Gemma for the approach—but JetMoE's contribution is in demonstrating that this curriculum is particularly effective under the constraints of a severely limited compute budget.

The field's default training strategy, inherited from the GPT-3 and Llama lineages, is to train on a single, carefully mixed data distribution with a cosine learning rate decay over the full training duration. This approach is robust but potentially inefficient: the model spends its entire training budget processing a fixed mixture, meaning that the last tokens seen during training—which receive the lowest learning rate and thus have the least impact on parameters—are drawn from the same distribution as the first tokens. The two-phase curriculum inverts this logic. The high learning rate in Phase 1 allows the model to extract maximal signal from the broad but noisy web data without overfitting. The Phase 2 switch to high-quality data with a decaying learning rate effectively performs a controlled "fine-tuning" at the end of pretraining, using the highest-quality signal to refine the model's knowledge when parameter updates are small and precise.

The intellectual contribution is the recognition that under a fixed, small compute budget, data quality is not a uniform commodity to be averaged but a scheduling resource to be deployed strategically. The paper's Phase 2 mixture (detailed in Table 2) allocates nearly 30% of the data to high-quality supervised and synthetic sources (OpenHermes, UltraTextbooks, TemplateGSM, Magicoder, etc.) that would typically be reserved for post-training instruction tuning. By pulling these datasets into the pretraining phase, the model benefits from instruction-following patterns and structured reasoning examples during the foundational learning stages, potentially explaining why the model achieves strong downstream performance despite its limited token budget.

This is an incremental advance in training methodology—it adapts a known technique (two-phase training) to the specific constraints of budget-limited training and provides a fully documented, reproducible recipe. Its significance lies in the practical demonstration that this approach works at scale with open data, not in theoretical novelty. The paper does not compare against a single-phase baseline, so the contribution is again about the demonstrated effectiveness of the combined recipe rather than a causal claim about the curriculum's necessity. The takeaway for practitioners is: when training on a tight budget, don't just pick one data mixture and one LR schedule—deploy your best data at the end, when you're making the most precise adjustments to the model.

5. Experimental Analysis

Evaluation Methodology

Dataset. The primary evaluation is conducted across the six benchmarks comprising the OpenLLM Leaderboard (ARC-challenge, Hellaswag, MMLU, TruthfulQA, WinoGrande, GSM8k), with additional evaluations on code generation benchmarks (HumanEval and MBPP) and the chat-oriented MT-Bench. The paper does not specify the number of test examples for each benchmark, though these are standard public benchmarks with well-known sizes (e.g., GSM8k contains 1,319 test problems, MMLU contains ~14,000 questions across 57 subjects, HumanEval contains 164 programming problems). The evaluation methodology follows the same protocol as the OpenLLM leaderboard to ensure direct comparability with other models, and uses the BigCode Evaluation Harness for code benchmarks.
Base model(s). The paper evaluates a single base model, JetMoE-8B, trained from scratch as described in Sections 3 and 4. Two variants are evaluated: the base pretrained model (JetMoE-8B) on the OpenLLM leaderboard and code benchmarks, and the aligned chat model (JetMoE-8B-Chat) on MT-Bench. There is no model family comparison—the paper does not train or evaluate smaller/larger JetMoE variants, making this a single-model study rather than a scaling analysis.
Metrics. For the OpenLLM leaderboard benchmarks, the primary metric is accuracy (percentage of correct answers), computed using the standard evaluation harness for each task. For code generation, the metric is pass@1 (the fraction of problems for which a single greedy-decoded solution passes all test cases), reported as a mean across HumanEval and MBPP. For MT-Bench, the metric is a scalar score on a 1–10 scale assigned by a strong LLM judge (gpt-4-0613) evaluating multi-turn conversation quality across eight categories (Writing, Roleplay, Extraction, STEM, Humanities, Reasoning, Math, Coding), with the final score being the average across all categories. The MT-Bench evaluation uses task-specific temperatures: 0.7 for Writing and Roleplay (high creativity), 0.0 for Extraction, Math, Coding, and Reasoning (precise), and 0.1 for STEM and Humanities (slight variability).
Baselines. The paper compares against three external open-source models in Table 3: Llama2-7B (Touvron et al., 2023), DeepseekMoE-16B (Dai et al., 2024), and Gemma-2B (Team et al., 2024). For the chat evaluation in Table 4, baselines include Llama2-7B-Chat, Llama2-13B-Chat, Vicuna-13B-v1.5, Zephyr-7B-beta, OpenChat-3.5, Starling-LM-7B-alpha, and GPT-3.5-turbo. The paper does not train its own dense baseline (e.g., a 2B-parameter dense model with equivalent training compute to JetMoE-8B), nor does it include the open-source MoE baseline OpenMoE in the comparison tables despite citing it as motivation—a notable omission given that OpenMoE's underperformance was a key part of the paper's stated motivation.
Generation budget / compute accounting. For the pretrained model evaluation, all comparisons are based on training budget: JetMoE-8B was trained on 1.25T tokens using 30,000 H100 GPU hours at a cost of approximately $0.1 million. The baseline models' training costs are not directly reported, but Llama2-7B was trained on 2T tokens (1.6× more data than JetMoE-8B) and DeepseekMoE-16B on 2T tokens (same ratio). At inference time, JetMoE-8B activates 2.2B parameters per token (Table 1), compared to 7B for Llama2-7B and 2B for Gemma-2B, giving JetMoE an inference efficiency advantage of roughly 70% versus Llama2-7B. There is no FLOPs-matched training-to-inference tradeoff analysis analogous to the compute-optimal test-time scaling framework in the reference example—the paper's cost claim is about training economics, not inference-time resource allocation.
Cross-validation / statistical protocol. The paper performs no cross-validation or statistical significance testing. All evaluation results in Tables 3 and 4 are reported as single-point estimates from the standard benchmark evaluation harnesses. There are no error bars, confidence intervals, or multiple evaluation runs reported. For a paper whose central claim is about cost-effectiveness, the absence of any statistical protocol means we cannot assess whether JetMoE-8B's 1.9-point advantage over Llama2-7B on the OpenLLM Leaderboard average (53.0 vs. 51.0) is reliable or falls within the expected variance of these benchmarks. The MT-Bench evaluation uses a single judge model (gpt-4-0613) without reporting inter-judge agreement or multiple judge runs, a known source of variance in LLM-as-judge evaluations.

Main Quantitative Results

OpenLLM Leaderboard and Code Benchmarks

The headline result in Table 3 is that JetMoE-8B achieves an average OpenLLM Leaderboard score of 53.0, outperforming Llama2-7B (51.0), DeepseekMoE-16B (51.1), and Gemma-2B (46.4). This is the basis for the paper's central claim that "JetMoE-8B outperforms the Llama2-7B model." However, the task-level breakdown reveals a more nuanced picture:

Where JetMoE wins decisively: GSM8k (27.8 vs. 14.5 for Llama2-7B, a 92% relative improvement), TruthfulQA (41.7 vs. 38.8), and MMLU (49.2 vs. 46.9). The GSM8k result is strikingly strong—this is a mathematical reasoning benchmark where Llama2-7B is notably weak, and JetMoE's 27.8 approaches the level of much larger models. Hellaswag is also a clear win (80.5 vs. 78.6).
Where JetMoE loses: ARC-challenge (48.7 vs. 53.1 for Llama2-7B, trailing by 4.4 points) and WinoGrande (70.2 vs. 74.0). These are commonsense reasoning benchmarks that require broad world knowledge, and JetMoE's underperformance here—despite having more total parameters—raises questions about whether the 1.25T token budget (vs. Llama2's 2T) or the open-source data mixture leaves gaps in commonsense coverage.
Where JetMoE is roughly tied: The remaining benchmarks show JetMoE competitive but not dominant. The average advantage of 2.0 points over Llama2-7B is driven disproportionately by the GSM8k and TruthfulQA results; removing GSM8k from the average would narrow the gap substantially.

On code generation, JetMoE-8B achieves a MBPP pass@1 of 34.2, matching DeepseekMoE-16B (34.0) and substantially exceeding Llama2-7B (20.8) and Gemma-2B (28.0). However, on HumanEval, JetMoE-8B achieves only 14.6, trailing DeepseekMoE-16B (25.0), Gemma-2B (24.4), and barely exceeding Llama2-7B (12.8). This asymmetry—strong on MBPP (simpler, more templated Python problems) but weak on HumanEval (more diverse algorithmic challenges)—is consistent across several models in the table, but the 19.6-point gap between MBPP and HumanEval for JetMoE is the largest among the compared models. The paper does not discuss this discrepancy, which could indicate overfitting to the types of solutions found in the MBPP training distribution (which overlaps with the code datasets in Phase 2 training) or a failure to generalize to novel algorithmic problems.

Comparing against DeepseekMoE-16B is particularly instructive for understanding JetMoE's architectural claims. DeepseekMoE-16B is a much larger model (16B total, 2.8B active) trained on 2T tokens, with a more sophisticated MoE design that the paper cites as inspiration. JetMoE-8B (8B total, 2.2B active, 1.25T tokens) matches or exceeds DeepseekMoE on nearly every benchmark except ARC-challenge (48.7 vs. 53.2) and HumanEval (14.6 vs. 25.0). This is genuinely impressive—a model with half the total parameters, 60% of the training data, and a simpler architecture is competitive. However, this comparison partially undermines the paper's narrative about MoA being crucial: DeepseekMoE uses dense attention and still achieves comparable overall performance, suggesting that the efficiency gains may come more from the data curriculum and training recipe than from the attention sparsity innovation.

MT-Bench Chat Evaluation

The chat evaluation in Table 4 and Figure 3 shows JetMoE-8B-Chat achieving an MT-Bench score of 6.48, surpassing Llama2-13B-Chat (6.27) and Llama2-7B-Chat (6.03). This is the basis for the paper's claim that "JetMoE-8B-Chat surpasses the Llama2-13B-Chat model."

The two-turn score (6.69) exceeds the single-turn score (6.16), which is unusual—most models in the comparison show the opposite pattern, with multi-turn conversations being harder than single-turn interactions. The paper does not comment on this anomaly, but it could indicate that the dSFT and dDPO alignment process, which uses multi-turn datasets like UltraChat 200k and Capybara, was particularly effective at training the model for sustained dialogue.

The radar plot in Figure 3 reveals clear strengths and weaknesses relative to GPT-3.5-turbo (the strongest model in the comparison at 8.14). JetMoE-8B-Chat is notably weak in Coding and Extraction compared to GPT-3.5-turbo, with large gaps visible in the radar plot. The paper attributes this to "smaller model size leading to suboptimal reasoning capability in these tasks." However, since JetMoE-8B-Chat outperforms Llama2-13B-Chat overall—and Llama2-13B has more active parameters (13B vs. 2.2B)—the weakness cannot simply be attributed to parameter count. It more likely reflects the specific data mixture: the code SFT data in Phase 2 (Magicoder, Evol-Code Alpaca, Code-290k-ShareGPT) may not adequately cover the diversity of coding tasks tested in MT-Bench, or the dSFT/dDPO alignment may have been optimized toward certain interaction patterns at the expense of code-related capabilities.

The comparison with Zephyr-7B-beta (7.34), Starling-LM-7B-alpha (8.09), and OpenChat-3.5 (7.49) in Table 4 reveals that JetMoE-8B-Chat underperforms these strong 7B-scale chat models by a substantial margin. These models use similar alignment techniques (dSFT and dDPO, particularly Zephyr which pioneered the approach) but are built on the dense Mistral-7B base model. This comparison suggests that while JetMoE-8B's pretrained model is competitive with Llama2-7B, the aligned chat variant does not close the gap with the best 7B-scale chat models, despite having 8B total parameters. The paper does not discuss this limitation, focusing instead on the Llama2 comparison.

Inference Efficiency

While not presented as a separate results section, the activation parameter comparison in Table 3 is a form of efficiency result. JetMoE-8B activates 2.2B parameters per token versus 7B for Llama2-7B and Llama2-13B (presumably 13B, though the chat variant's activation count is not specified), giving it a 3.2× advantage in per-token active parameters. This means the benchmark results in Table 3 are achieved with approximately 70% less inference computation than Llama2-7B, which the paper highlights as a key advantage. However, this comparison does not account for the overhead of the router computation, the shared key/value projection, or the load balancing machinery—components that add a constant cost regardless of which experts are activated. The effective speedup in practice would be less than the 70% parameter-count reduction would suggest, and the paper provides no latency or throughput measurements to quantify the actual inference-time savings.

Ablation Studies and Robustness Checks

No ablation studies were performed. The paper explicitly states this in Section 7:

"Due to the limited $100k budget, we can not afford any ablation study for the model architecture."

This is a significant limitation for a paper whose stated purpose is to provide a reproducible recipe and advance open MoE research. Without ablations, we cannot answer any of the following questions, all of which are central to evaluating the paper's claimed contributions:

Does MoA (sparse attention) improve performance over a standard dense-attention MoE? This is the paper's main architectural innovation, yet there is no experiment comparing JetMoE against an identical model where only the FFN layers use SMoE and attention is dense. The DeepseekMoE-16B comparison in Table 3 provides a weak form of evidence (DeepseekMoE uses dense attention and achieves similar overall performance), but the models differ in total parameters, training data, and data mixture, making causal attribution impossible.
Does the shared KV projection help? The paper argues that sharing key and value matrices across attention experts "improves training stability," but there is no training stability comparison against an architecture with per-expert KV projections, nor any measurement of whether training was in fact more stable (e.g., loss curves, gradient norms, or router collapse incidents).
Does the two-phase data curriculum matter? The paper presents the Phase 1/Phase 2 split as a key design choice inspired by miniCPM and Gemma, but there is no comparison against a model trained on the same data mixture with a single-phase cosine decay schedule. The paper cannot distinguish whether the gains come from the two-phase approach or from the specific data mixtures, which differ in composition (Phase 2 includes much more synthetic and instruction data regardless of scheduling).
Does the WSD learning rate schedule matter? The Warmup-Stable-Decay schedule is another design choice credited to miniCPM, but there is no comparison against a standard cosine decay on the same data.
Do the load balancing and z-loss weights matter? The values α = 0.01 and β = 0.001 are taken directly from ST-MoE without tuning. There is no evidence that these are appropriate for the JetMoE architecture, nor any sensitivity analysis showing that performance is robust to these choices.
What is the contribution of each data source? The Phase 2 data mixture (Table 2) includes nearly 30% synthetic and instruction-tuning data. There is no experiment comparing against a model trained on web data alone, making it impossible to determine whether JetMoE's performance comes from the MoE architecture or from training on higher-quality data than typical pretraining runs.

The absence of ablations is understandable given the budget constraints—running even a single ablation would double the already-substantial cost. However, it fundamentally limits the paper's contribution from "we demonstrate this approach works and why" to "we demonstrate this combination works, but we don't know which parts matter." The paper is honest about this limitation, but honesty does not remedy the gap in causal understanding.

The only comparisons that function as weak ablations are the external model comparisons in Tables 3 and 4. These allow some qualitative inference:

MoE vs. dense architecture: Comparing JetMoE-8B (MoE, 2.2B active) against Llama2-7B (dense, 7B active) and Gemma-2B (dense, 2B active) provides suggestive evidence that the MoE architecture is beneficial—JetMoE beats both with similar or fewer active parameters. However, these models differ in training data quality/quantity, tokenizer, and training recipe, so this is not a controlled comparison.
Two-phase data curriculum: Comparing JetMoE-8B against OpenLLaMA (which the paper cites in Section 1 as a weak dense model) is suggestive—OpenLLaMA was trained on open data but without the Phase 2 high-quality mixture and WSD schedule, and JetMoE substantially outperforms it. However, OpenLLaMA is not included in the evaluation tables, so this comparison is only implied by the motivation section.
Alignment effectiveness: Comparing JetMoE-8B-Chat (dSFT + dDPO, 60 H100 hours) against Llama2-7B-Chat (RLHF, cost undisclosed) shows that the simpler alignment pipeline achieves better MT-Bench scores, suggesting the dSFT/dDPO approach is effective. But the base models differ, so this is not a clean ablation of alignment methods.

The ReST-like experiment is absent. Section 5.2 discusses online and iterative DPO as alternatives but does not implement or evaluate them. The paper would benefit from even a small-scale comparison between offline dDPO (as used) and a simple iterative variant, to test whether the distribution shift concern raised in Section 5.2 affects JetMoE in practice.

No training stability metrics are reported. The paper introduces z-loss specifically for training stability but provides no loss curves, router utilization plots, or expert load distributions over the course of training. For a paper whose contributions include a novel attention sparsity mechanism and specific stability techniques, the complete absence of training dynamics data is a significant gap. The reader cannot assess whether the auxiliary losses worked as intended, whether the router remained balanced, or whether training was in fact stable.

Critical Assessment

The experiments, while demonstrating JetMoE-8B's competitive performance, provide only partial support for the paper's central claims, and the missing ablations create fundamental ambiguity about which design choices are responsible for the results.

Claim: "LLM training can be much more cost-effective than generally thought." The experiments support this claim as an existence proof but not as a systematic demonstration. JetMoE-8B was indeed trained for ~$0.1 million and achieves competitive performance. However, the claim's force depends on what baseline cost expectation is being challenged. If the implicit assumption is that training a Llama2-competitive model costs several million dollars, then a single successful run at 10× lower cost is strong evidence that the assumption was wrong. But the paper does not establish what the field's "generally thought" cost expectation actually is—it treats this as common knowledge without citation. Additionally, the $0.1 million figure is not independently verifiable; it depends on the specific GPU pricing, utilization efficiency, and infrastructure deals the authors had access to. A different research group attempting to replicate the training might face higher costs. The claim would be more convincing with: (a) a specific citation for the cost assumption being challenged, (b) a breakdown of the $0.1 million calculation, and (c) demonstration that the cost includes failed runs, debugging, and infrastructure setup time, not just the final successful training.

Claim: "JetMoE-8B outperforms the Llama2-7B model." This claim is supported on average (53.0 vs. 51.0 OpenLLM average in Table 3) but the average masks important task-level variation. JetMoE loses on ARC-challenge (-4.4 points) and WinoGrande (-3.8 points), wins on GSM8k (+13.3 points) and TruthfulQA (+2.9 points), and is essentially tied on Hellaswag (+1.9 points). The "outperforms" claim is true for the aggregate metric but would be false if restricted to commonsense reasoning tasks. More importantly, the paper does not establish the statistical significance of the 2.0-point average gap. With six benchmarks and no reported variance, it is impossible to determine whether running the evaluation with a different random seed or on a different test set split would reverse the ordering. The Llama2 results are presumably taken from the OpenLLM Leaderboard as reported, but those scores themselves have measurement uncertainty that is not propagated into the comparison. Given the narrow margin, the claim should be qualified as "competitive with" rather than "outperforms."

Claim: "JetMoE-8B-Chat surpasses the Llama2-13B-Chat model." This claim is supported but fragile. The MT-Bench margin is 6.48 vs. 6.27 (Table 4), a difference of 0.21 points. MT-Bench scores from a single LLM judge have well-documented variance (Zheng et al., 2023, the very paper that introduced MT-Bench, reports that GPT-4 judgments have a ~70-80% agreement rate with human preferences, not perfect reliability). A 0.21-point difference on an 8.14-point range is a 2.6% relative gap—almost certainly within the noise floor of single-judge LLM evaluation. Additionally, JetMoE-8B-Chat trails the best 7B-scale chat models (Starling, Zephyr, OpenChat) by substantial margins (1.01–1.61 points). The "surpasses Llama2-13B-Chat" framing is technically true but omits the context that JetMoE-8B-Chat is not competitive with the state-of-the-art in its own size class. A more honest framing would be "marginally exceeds Llama2-13B-Chat on MT-Bench while substantially underperforming contemporaneous 7B chat models."

What would strengthen the paper:

A single controlled ablation. Even one experiment—e.g., training a JetMoE variant with dense attention but all other components identical—would provide a causal anchor for the architecture claims. The 30,000 H100 hours budget was apparently exhausted, but a scaled-down experiment (e.g., training to 500B tokens instead of 1.25T) would still provide directional evidence.
Training dynamics data. Loss curves (LM loss and auxiliary losses), router utilization over time, and expert load distributions would allow readers to assess whether the proposed stability mechanisms (z-loss, frequency-based load balancing) functioned as intended, and whether the claimed advantages of shared KV projections for stability are reflected in training behavior. These are inexpensive to log during training.
A fair MoE baseline at comparable compute. Including OpenMoE in Table 3 would directly test whether JetMoE's architecture is responsible for the improved performance over prior open MoE attempts, or whether the data mixture and training recipe account for the difference. OpenMoE was designed for the same purpose (open MoE research) and trained on open data—excluding it from the comparison is a significant omission given that OpenMoE's weakness was a key motivation.
Inference throughput measurements. The claim of 70% inference computation reduction relative to Llama2-7B is purely based on active parameter counts. Actual throughput depends on the efficiency of the sparse computation kernels (MegaBlocks), the overhead of the router, and the KV cache management for the shared key/value projections. Latency and throughput measurements on standard hardware would transform the efficiency claim from a theoretical estimate to a demonstrated advantage.
A cost breakdown. The $0.1 million figure is stated without any line-item detail. How many GPU hours were spent on the final successful run versus debugging and infrastructure setup? What was the per-GPU-hour cost? Were any costs (e.g., data preprocessing, infrastructure support) not included? For a paper whose primary contribution is a cost-effectiveness demonstration, this opacity undermines the reproducibility the paper otherwise champions.

The contribution that is genuinely demonstrated: JetMoE-8B shows that a model with 8B total / 2.2B active parameters, trained on 1.25T tokens of carefully composed open-source data, can achieve performance broadly comparable to Llama2-7B (a 7B-parameter dense model trained on 2T tokens of proprietary data) at what the authors claim is substantially lower cost. This is a useful existence proof and a valuable contribution of a fully-documented recipe. The experiments cleanly support this narrower claim.

The contribution that is not demonstrated: That the MoA architecture, the shared KV projections, the two-phase curriculum, the WSD schedule, or any other individual design choice is responsible for the results. The paper's framing suggests these are the innovations that enable the cost reduction, but the experiments provide no causal evidence for any of them. Given that the paper explicitly states it cannot run ablations, the contribution is best understood as a recipe demonstration rather than a validated design principle—a distinction the paper itself acknowledges but that may be lost in the broader narrative.

6. Limitations and Trade-offs

6.1 Complete Absence of Ablation Studies Leaves Causal Attribution Unknown

The assumption or constraint. The paper's central architectural claim is that extending sparsity to attention layers (MoA) with shared key/value projections improves efficiency and enables balanced scaling. However, the paper explicitly acknowledges in Section 7:

"Due to the limited $100k budget, we can not afford any ablation study for the model architecture. The hyperparameters and data mixtures are also handpicked based on the empirical results from previous works."

This is not a minor caveat—it means that none of the paper's proposed innovations are individually validated. Every design choice (MoA architecture, shared KV projections, two-phase data curriculum, WSD learning rate schedule, load balancing coefficients, specific data mixtures) is evaluated only as part of a single combined recipe. There is no experiment isolating any component's contribution.

The consequence. A practitioner reading this paper cannot determine which design choices are necessary versus incidental. The MoA architecture—the paper's signature innovation—could be irrelevant to the model's performance; the same results might be achievable with dense attention and more FFN experts given the same training recipe and data. The two-phase data curriculum might be responsible for most of the gains, with the sparse attention contributing little. The shared KV projection might be unnecessary for stability if z-loss alone suffices. This creates a fundamental reproducibility problem: a group attempting to build on JetMoE cannot focus their efforts on the components that matter most, and a group that modifies one component (e.g., switching to dense attention to simplify implementation) cannot predict whether performance will degrade, stay constant, or improve.

More subtly, the absence of ablations makes the paper's cost-efficiency claim non-actionable. The paper demonstrates that this specific configuration achieves competitive performance at low cost, but provides no guidance on how to trade off components against budget. Would training with dense attention but the same data curriculum cost 20% more and perform better? Would removing the Phase 2 high-quality data still yield a competitive model at even lower cost? These are precisely the questions a practitioner needs answered when deciding whether to adopt the recipe.

What evidence exists in the paper. The only evidence that functions as even a weak ablation is the external model comparison in Table 3. DeepseekMoE-16B uses dense attention and achieves comparable overall performance (51.1 vs. JetMoE's 53.0 OpenLLM average), which undermines rather than supports the necessity of MoA—a model without sparse attention can perform similarly. However, DeepseekMoE differs in total parameters (16B vs. 8B), training data (2T vs. 1.25T tokens), data mixture (proprietary vs. open), and many other dimensions, so this is not a controlled comparison. The paper provides no loss curves, no router utilization data, and no measurements of whether the shared KV projection actually improved training stability as claimed.

Mitigation status. The paper acknowledges this limitation transparently in Section 7 and suggests that "in the future, it would be interesting to further study the actual contribution of different components to the final results." This is an honest admission but provides no immediate solution for the reader trying to evaluate the paper's claims. The limitation is structural—it cannot be mitigated within the current work—and means the paper's contribution is best understood as a recipe demonstration rather than a set of validated design principles.

6.2 The Single Budget-Constrained Training Run Provides No Evidence of Scaling Behavior

The assumption or constraint. All results in the paper come from a single JetMoE-8B model trained on exactly 1.25T tokens with exactly 8 experts per layer and top-2 routing. The paper implicitly extrapolates from this single data point to make claims about the general effectiveness of the approach: "LLM training can be much more cost-effective than generally thought." However, the paper provides no evidence about how performance varies with model scale, training tokens, or number of experts. This is a single-model study, not a scaling analysis.

The consequence. A practitioner cannot answer basic planning questions from this paper: What happens if I train JetMoE on 2T tokens instead of 1.25T—does performance continue to improve, plateau, or degrade? What if I use 16 experts instead of 8? What if I build a JetMoE-3B variant with a proportionally smaller budget—does the cost-effectiveness advantage hold at smaller scales, or is there a threshold below which the MoE overhead dominates? The paper's cost-effectiveness claim is tied to one specific operating point, and there is no evidence that this point represents an optimum or that the cost advantage generalizes.

This is particularly important because the paper's headline claim—$0.1 million achieves Llama2-7B performance—relies on a comparison where the models are not matched on training compute. JetMoE-8B was trained on 1.25T tokens, Llama2-7B on 2T tokens (1.6× more). If JetMoE were trained on a matched 2T tokens, its performance might improve further (strengthening the cost claim) or might plateau (weakening it relative to the per-token efficiency framing). Without scaling data, the reader cannot extrapolate.

What evidence exists in the paper. None. There are no experiments varying model size, expert count, training tokens, or top-k. The paper does not even include training loss curves that would indicate whether the model was still improving when training stopped at 1.25T tokens—did the loss plateau, or was it still declining? This information would at least hint at whether additional training would help. The comparisons in Table 3 and Table 4 are all against external models at different scales, which provides point estimates of relative performance but no within-family scaling trends.

Mitigation status. Not addressed. The paper does not discuss this as a limitation or suggest future scaling studies. The WSD learning rate schedule (Section 4.2) is described with specific phase lengths (10B token warmup, 250B token decay), implying the authors believed 1.25T tokens was an appropriate total, but there is no justification for this choice beyond reference to prior work (miniCPM). The inference efficiency claim—70% less computation than Llama2-7B—depends on the 8-expert, top-2 configuration, but there is no evidence that this configuration is near-optimal for the accuracy-efficiency tradeoff curve.

6.3 Data Mixture and Curriculum Are Confounded With Architecture; Neither Is Isolated

The assumption or constraint. JetMoE-8B's training recipe includes both a novel architecture (MoA + SMoE) and a carefully curated, two-phase data mixture that differs substantially from standard pretraining recipes. Phase 2 (Table 2) allocates approximately 30% of the data to high-quality synthetic and instruction-tuning datasets—OpenHermes, UltraTextbooks, Magicoder, TemplateGSM, and others—that would typically be reserved for post-training fine-tuning, not included in the main pretraining corpus. The paper presents this as part of the efficiency recipe but provides no experiment that separates the effect of the data from the effect of the architecture.

The consequence. This creates a fundamental confound: we cannot distinguish whether JetMoE-8B outperforms Llama2-7B because of its architecture or because it was trained on better data. Llama2-7B was trained on a proprietary, undisclosed mixture dominated by web data, with instruction tuning applied only after pretraining as a separate phase. JetMoE-8B integrates instruction-style data directly into the pretraining curriculum during Phase 2, effectively performing a form of instruction-tuning during pretraining. If a dense model with the same parameter count and training budget were trained on JetMoE's exact data mixture, it might match or exceed JetMoE-8B's performance—and we would have no way to know.

This confound also affects the paper's open-source narrative. One of the paper's stated motivations is that commercial MoE models are "not truly open-sourced as the training recipes are not published and may contain proprietary datasets." JetMoE addresses the recipe transparency but introduces a new ambiguity: the recipe's success may be attributable to the unusually high proportion of synthetic and instruction data, which the paper provides citations for but does not ablate. A practitioner trying to reproduce the results might reasonably ask: if I use the same data mixture but a dense architecture, do I get similar performance for similar cost?

What evidence exists in the paper. The Phase 2 data mixture is meticulously documented in Table 2 with exact percentages (e.g., RefinedWeb 39.8%, Starcoder 19.6%, UltraTextbooks 4.8%, OpenHermes and related synthetic data ~7.3%). However, there is no comparison against a model trained on (a) the same architecture but a standard single-phase web-only mixture, or (b) the same data mixture but a dense architecture. The paper cites miniCPM and Gemma as inspiration for the two-phase approach, noting they also use this strategy, but does not establish whether these models also confound data quality with architecture.

Mitigation status. The paper partially acknowledges this indirectly: "given the limited computing budget available, our data mixture might not be ideal. However, it serves as a good starting point for training JetMoE-8B and can be further optimized in future iterations." This admits the mixture is not claimed to be optimal but does not address the confounding issue. The paper does not suggest future work to disentangle architecture from data, focusing instead on architectural ablations as the priority for future work.

6.4 Evaluation Statistical Reliability Is Unquantified; Narrow Margins May Not Be Robust

The assumption or constraint. All evaluation results in Tables 3 and 4 are reported as single-point estimates with no confidence intervals, no error bars, no cross-validation runs, and no multiple evaluation seeds. The paper uses the standard OpenLLM Leaderboard evaluation harness and reports scores to one decimal place, but provides no information about the statistical reliability of these numbers—either for JetMoE-8B or for the baseline models, which are taken from the leaderboard as reported. The MT-Bench evaluation uses a single LLM judge (gpt-4-0613) without reporting inter-judge agreement or multiple judge runs.

The consequence. Several of the paper's headline claims rest on margins that are small enough to potentially fall within evaluation noise. JetMoE-8B's 2.0-point OpenLLM Leaderboard advantage over Llama2-7B (53.0 vs. 51.0) is a 3.9% relative difference, averaged across six benchmarks. On individual benchmarks, JetMoE loses to Llama2 on ARC-challenge (48.7 vs. 53.1) and WinoGrande (70.2 vs. 74.0), making the average gap sensitive to which benchmarks are included and their relative weighting. The MT-Bench advantage over Llama2-13B-Chat (6.48 vs. 6.27) is 0.21 points on a ~8-point scale—a 2.6% relative gap that is almost certainly within the noise floor of single-judge LLM evaluation, given that Zheng et al. (2023) report GPT-4 judgments have approximately 70–80% agreement with human preferences, not perfect reliability.

A practitioner deciding between JetMoE-8B and Llama2-7B for a specific application cannot determine from this paper whether the observed differences are reliable or would reverse with a different evaluation setup. The paper provides no information about whether JetMoE-8B was evaluated multiple times, whether the baseline scores include their own variance, or what the expected test-retest reliability is for any of these benchmarks.

What evidence exists in the paper. The single-point estimates in Tables 3 and 4. The paper states that evaluation follows "the same evaluation methodology as in the OpenLLM leaderboard to be comparable to other models," which implies single-run evaluation since the leaderboard typically reports point estimates. The absence of error bars is not discussed. The MT-Bench methodology section specifies temperatures for different task categories (0.7 for Writing/Roleplay, 0.0 for Extraction/Math/Coding/Reasoning, 0.1 for STEM/Humanities) but does not specify how many evaluation runs were performed or how scores were aggregated across turns and categories beyond "average."

Mitigation status. Not addressed. The paper does not discuss statistical reliability as a concern, does not report variance, and does not suggest future work on robust evaluation. This is a common limitation in LLM benchmark reporting, but it is particularly consequential here because the paper's central claims depend on relatively narrow margins and because the paper's contribution is explicitly about cost-effectiveness—a claim that requires quantifying not just point estimates of performance but also the uncertainty around those estimates to make informed resource-allocation decisions.

6.5 The $0.1 Million Cost Claim Is Not Independently Verifiable or Replicable

The assumption or constraint. The paper's most attention-grabbing claim is that JetMoE-8B was trained with a "$100k budget" using "30,000 H100 GPU hours." This figure is central to the paper's narrative about democratizing LLM research, but it is presented without any breakdown or verifiable accounting. The paper does not specify: the cost per GPU hour, whether the 30,000 hours includes only the final successful training run or also the inevitable debugging, failed runs, hyperparameter exploration, and infrastructure setup; whether infrastructure costs (networking, storage, data preprocessing) are included; or whether the authors had access to subsidized or institutional pricing that may not be available to other researchers.

The consequence. A research group attempting to replicate JetMoE-8B from scratch cannot budget based on the $0.1 million figure without understanding what it includes. If the figure represents only the final training run and the group must spend additional thousands of GPU hours on debugging, hyperparameter search, and infrastructure setup, their actual cost may be 2–3× higher. If the per-GPU-hour rate reflects institutional or cloud-provider discounts unavailable to smaller groups, the effective cost at market rates could be substantially higher. This undermines the paper's core claim about accessibility: a recipe is not truly reproducible if the reported cost is not achievable by the target audience.

More subtly, the cost figure is used in the abstract and introduction as evidence for the claim that "LLM training can be much more cost-effective than generally thought," but the paper never establishes what "generally thought" means—there is no citation for the cost of Llama2-7B training, no estimate of what a comparable dense model would cost at market rates, and no baseline cost expectation. The $0.1 million figure is presented as surprising, but the reader has no way to verify whether it is actually lower than what a straightforward scaling analysis would predict for a 1.25T-token training run on 96 H100s.

What evidence exists in the paper. Section 4.1 describes the infrastructure: "a cluster containing 12 nodes and 96 H100s. Inside each node, gpus are connected via NVLinks. Infiniband is used for fast communication between nodes." This provides the hardware setup but not the cost. The 30,000 H100 hours figure is stated in the abstract and again in the conclusion without further detail. The aligned model (Section 5.3) uses an additional 60 H100 GPU hours for the dSFT + dDPO process, suggesting the base model training consumed the vast majority of the 30,000 hours. No acknowledgments section mentions special pricing, donated compute, or institutional subsidies that would affect cost.

Mitigation status. Not addressed. The paper does not discuss the cost calculation methodology, the scope of costs included, or the generalizability of the per-GPU-hour rate. The "existence proof" framing (that the authors managed to do it for this amount) is logically distinct from the "reproducible budget" framing (that others can expect to do it for this amount), and the paper conflates them. A simple cost breakdown in an appendix would substantially strengthen the reproducibility claim without requiring additional experiments.

6.6 MoA and SMoE Inference Overhead Is Not Measured; Actual Speedup Remains Unknown

The assumption or constraint. The paper claims that JetMoE-8B "reduc[es] inference computation by about 70% compared to Llama2-7B" based purely on active parameter counts: JetMoE activates 2.2B parameters per token versus Llama2-7B's 7B. This is a theoretical upper bound on computational savings that assumes zero overhead from the sparse computation machinery. In practice, the inference cost includes the router computation (a matrix multiply that runs for every token regardless of which experts are selected), the shared key/value projection (also always-on), the load balancing bookkeeping, the kernel launch overhead for sparse expert dispatch, and the memory bandwidth cost of loading expert parameters that may not be contiguous in memory. None of these overheads are measured.

The consequence. The actual inference speedup on standard hardware could be substantially less than the 70% parameter-count reduction would suggest. For a practitioner deploying JetMoE-8B in a latency-sensitive application, the theoretical parameter advantage is irrelevant if the sparse kernels, router overhead, or irregular memory access patterns negate the savings. Different hardware (e.g., consumer GPUs vs. datacenter H100s, CPU inference, edge devices) may show different overhead characteristics, and the paper provides no guidance about which deployment scenarios benefit most from the architecture.

This limitation is particularly important because routing and sparse operations are known bottlenecks in MoE inference. The MegaBlocks library (Gale et al., 2023) that JetMoE uses was specifically designed to address these inefficiencies, suggesting the authors were aware of the challenge, but they provide no benchmark results demonstrating that MegaBlocks achieves close-to-theoretical throughput on JetMoE's specific architecture. The shared KV projection in the MoA design is argued to improve efficiency by avoiding per-expert key/value computation, but the actual throughput benefit relative to a per-expert KV design is unmeasured.

What evidence exists in the paper. None. There are no throughput measurements, no latency benchmarks, no comparisons of tokens-per-second between JetMoE-8B and Llama2-7B on any hardware configuration. The efficiency claim is entirely based on active parameter counts. The paper does not discuss the computational cost of the router, the shared KV projection, or the load balancing mechanism, nor does it provide any profiling data from the training runs that would indicate the fraction of time spent on sparse dispatch versus expert computation.

Mitigation status. Not addressed. The paper makes the 70% claim in the abstract and introduction as a key advantage without qualifying it as a theoretical upper bound based on parameter counts. There is no suggestion of future work on inference benchmarking or deployment optimization. For a paper whose primary contribution is an efficiency argument, the absence of any real-system measurement of that efficiency is a significant gap between the claimed advantage and its demonstrated practical value.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper changes the landscape not by introducing a novel technique but by demonstrating that the assumed cost floor for competitive LLM training is substantially higher than necessary. The shift is less about any individual architectural innovation and more about the reframing of what is possible within a constrained academic or small-company budget. Prior to JetMoE, the field's implicit consensus—reinforced by the compute requirements reported for models like Llama2, Pythia, and even the open-source OpenMoE—was that training a model competitive with Llama2-7B required millions of dollars, proprietary data, and infrastructure only available within large corporate labs. JetMoE-8B challenges this not through argument but through a single, fully-documented existence proof: a model trained for under $0.1 million on open data that beats Llama2-7B on the OpenLLM Leaderboard (53.0 vs. 51.0 average) and, after alignment, surpasses Llama2-13B-Chat on MT-Bench (6.48 vs. 6.27).

The magnitude of this shift should be understood precisely. This is not a paradigm shift—the paper introduces no new theory, no new training algorithm, and no architectural principle validated through ablation. It is, instead, a practical reframing of the resource barrier: a demonstration that careful engineering of existing components (SMoE, MoA, WSD scheduling, two-phase data curricula) can produce a competitive model at roughly an order of magnitude lower cost than the field implicitly assumed. The paper's value is as a lower bound on cost-effectiveness—an existence proof that invites optimization—rather than an upper bound on what the architecture can achieve.

The paper also partially reconciles a tension in the open-source MoE literature. OpenMoE (Xue et al., 2024) had demonstrated that simply releasing an MoE model with open weights does not guarantee competitive performance; OpenMoE performed only on par with weak dense models of similar active parameter count. This created a pessimistic narrative: perhaps making MoEs work well requires proprietary data, undisclosed training tricks, or scale beyond what open-source efforts can muster. JetMoE-8B falsifies this narrative. It demonstrates that an open-source MoE can compete with strong dense baselines, and it does so while activating only 2.2B parameters per token—fewer than Gemma-2B. The resolution comes not from a single breakthrough but from the accumulation of careful design choices: applying sparsity to both attention and FFN layers, using shared KV projections for stability, adopting a two-phase data curriculum that injects high-quality synthetic data during the learning rate decay, and implementing robust load balancing with both frequency-based and z-loss auxiliary objectives.

The paper redirects research attention in several ways. First, it makes data mixture engineering a first-class research concern for budget-constrained training. The Phase 2 curriculum—allocating ~30% of training data to synthetic instruction and textbook sources during the decay phase—is a concrete, replicable strategy that the paper shows can substantially boost downstream performance. This shifts attention away from pure architectural innovation and toward the interaction between data scheduling and learning rate dynamics. Second, it makes MoA (sparse attention) a credible design axis for MoE architectures, even though the paper does not ablate it. Prior to JetMoE, the default assumption in the MoE literature was that attention should remain dense—Switch Transformer, GLaM, Deepseek MoE, and Mixtral all follow this pattern. JetMoE demonstrates that sparse attention is at least not harmful to performance when combined with the right stabilizing techniques (shared KV, z-loss), opening a design space that was previously considered too unstable or inefficient to explore. Third, it establishes the WSD learning rate schedule as a practical tool for budget-limited training, following miniCPM, by demonstrating its effectiveness in a fully open reproduction rather than a proprietary training run.

What becomes less attractive as a research direction: pure architectural innovation without accompanying data curriculum design. The paper's results, while not disentangling architecture from data, strongly suggest that the data mixture and its scheduling are first-order contributors to performance. A research program that focuses exclusively on novel sparse architectures while using standard web-only pretraining data may be leaving substantial gains on the table. Similarly, the paper's underperformance relative to state-of-the-art 7B chat models (Zephyr-7B-beta at 7.34, Starling-LM-7B-alpha at 8.09 on MT-Bench, versus JetMoE-8B-Chat at 6.48) suggests that alignment quality, not just pretraining efficiency, is a bottleneck for open MoE models seeking to match the best dense chat models.

Follow-Up Research This Work Enables

A controlled ablation of the MoA architecture against dense attention at matched compute. The paper's signature architectural claim—that extending sparsity to attention layers via MoA with shared KV projections improves efficiency—is entirely unvalidated by ablation. A direct follow-up would train two models on identical data mixtures and schedules, one with JetMoE's full MoA + SMoE architecture and one with dense attention + SMoE (the standard MoE design pattern), matching total parameters and training tokens. The key measurement would be whether MoA improves performance per unit of training or inference compute, or whether the same gains can be achieved with dense attention and additional FFN experts. The DeepseekMoE-16B comparison in Table 3 (dense attention, 51.1 OpenLLM average vs. JetMoE's 53.0) is suggestive that MoA is not essential, but the models differ in too many dimensions for causal inference. A clean ablation would either validate MoA as a useful design principle or reveal it as incidental complexity. This experiment could be run at reduced scale (e.g., 500B tokens, 4 experts) to control cost while still providing directional evidence.

Scaling behavior of the JetMoE architecture: expert count, top-k, and training tokens. The paper provides a single operating point (8 experts, top-2, 1.25T tokens) with no evidence about how performance varies along these axes. A systematic follow-up would train JetMoE variants at a fixed parameter budget while varying the number of experts (4, 8, 16), the top-k (1, 2, 4), and the total training tokens (500B, 1T, 2T), measuring both perplexity and downstream benchmark performance. The critical question is whether the efficiency advantage over dense models holds across scales or is specific to the paper's configuration. Does the MoE overhead (router computation, load balancing, sparse dispatch) become proportionally smaller or larger as the model grows? Is there a threshold below which dense models are actually more efficient because the sparse machinery's overhead dominates? This would transform JetMoE from a single recipe into a scaling law that practitioners can use for capacity planning.

Disentangling data mixture from architecture via matched-training comparisons. The paper's Phase 2 data mixture—allocating ~30% of training tokens to synthetic instruction, textbook, and curated code/math data during the learning rate decay—is a significant departure from standard pretraining recipes. A crucial follow-up would train a dense model (e.g., a 2B-parameter Llama-style architecture) on JetMoE's exact data mixture and two-phase schedule, then compare its performance to both JetMoE-8B and to a standard Llama-2B trained on web-only data. This would isolate how much of JetMoE's performance comes from the MoE architecture versus the data curriculum. If the dense model approaches JetMoE's performance, it would suggest that data quality and scheduling—not sparse architecture—are the primary cost-efficiency levers, fundamentally reframing the paper's contribution. If the dense model substantially underperforms JetMoE, it would strengthen the architectural claims and justify the added complexity of MoE training.

Inference throughput benchmarking of JetMoE on diverse hardware. The paper's 70% computation reduction claim is purely theoretical, based on active parameter counts. A necessary follow-up would benchmark JetMoE-8B against Llama2-7B and Gemma-2B on real hardware—H100s, A100s, consumer GPUs (RTX 4090), and CPU inference—measuring tokens-per-second, memory bandwidth utilization, and latency at various batch sizes. Key measurements include: the overhead of the router computation as a fraction of total inference time, the efficiency of MegaBlocks sparse dispatch on different GPU architectures, the KV cache memory savings from shared key/value projections, and whether the irregular memory access patterns of expert dispatch cause throughput degradation on bandwidth-constrained hardware. This would convert the paper's theoretical efficiency claim into actionable deployment guidance: on which hardware, and under which serving configurations (batch size, sequence length), does JetMoE actually deliver its promised speedup?

Improving alignment quality to close the gap with state-of-the-art 7B chat models. JetMoE-8B-Chat achieves a respectable MT-Bench score of 6.48, beating Llama2-13B-Chat (6.27) but substantially trailing Zephyr-7B-beta (7.34), Starling-LM-7B-alpha (8.09), and OpenChat-3.5 (7.49)—all built on the dense Mistral-7B base. This gap is particularly visible in the Coding and Extraction dimensions of the radar plot (Figure 3). A targeted follow-up would investigate whether the alignment gap stems from the base model's capabilities (i.e., JetMoE-8B's pretrained coding ability is weaker than Mistral-7B's, as suggested by the HumanEval results in Table 3) or from the alignment procedure (the specific dSFT datasets, the dDPO hyperparameters, or the chat template). Experiments would include: (a) aligning Mistral-7B using JetMoE's exact dSFT/dDPO pipeline to isolate alignment quality, (b) augmenting the dSFT phase with additional high-quality code and extraction data, and (c) experimenting with iterative DPO variants (which Section 5.2 mentions but does not implement) to address the distribution shift problem.

Training a fully open MoE at the 1B active parameter scale to maximize accessibility. The paper's 30,000 H100 hours is impressive but still beyond the reach of most academic labs without substantial compute grants. A natural follow-up would apply the JetMoE recipe at a smaller scale—e.g., a 1B active / 4B total parameter model trained on 500B tokens—targeting a total cost of $10,000–20,000. This would test whether the efficiency advantages (MoA, two-phase curriculum, WSD schedule) are preserved when scaling down, or whether there is a minimum scale below which the sparse machinery overhead dominates. A successful result would make the JetMoE approach accessible to essentially any academic group with a handful of GPUs, truly democratizing MoE research. The experiment would also provide the first data point on the lower bound of JetMoE's scaling curve, enabling linear interpolation for budget planning.

Practical Applications and Downstream Use Cases

Academic LLM research on a consumer-grade compute budget. The paper's primary practical contribution is enabling academic groups to train and study competitive MoE models without corporate-scale infrastructure. A lab with access to a single 8×A100 node (available through many university compute programs or cloud credits) could, following the JetMoE recipe, train a smaller variant (e.g., 4B total / 1B active parameters) and conduct ablation experiments that are impossible with the full 8B model. The paper's exhaustive documentation of data sources (all public, all cited in Section 3), training infrastructure (Megatron + MegaBlocks, Section 4.1), and hyperparameters (Table 1) means a graduate student could begin a reproduction attempt without reverse-engineering undisclosed details. The 60 H100 hours for alignment (Section 5.3) is low enough that even the chat variant is accessible for iterative experimentation with different alignment strategies.

Cost-efficient deployment of capable models with reduced inference latency. JetMoE-8B's 2.2B active parameter count—fewer than Gemma-2B and roughly 30% of Llama2-7B—makes it suitable for deployment scenarios where inference cost or latency is the primary constraint, provided the actual throughput matches the theoretical parameter reduction (which remains unmeasured; see Limitations). Applications include on-device or edge deployment of a capable base model, batch inference for data processing pipelines where per-token cost dominates the total budget, and real-time interactive systems where the 70% parameter reduction could translate to proportionally lower latency. The model's strong GSM8k performance (27.8, Table 3) and competitive MMLU (49.2) suggest it is particularly well-suited for mathematical and knowledge-intensive applications despite its small active footprint.

Reproducible baseline for MoE architecture research. The paper provides the first fully open, competitive MoE recipe against which future architectural innovations can be compared. A researcher proposing a new routing mechanism, a new load balancing scheme, or a new sparse attention design can train a JetMoE-variant with their modification and compare directly against the published checkpoint, with the confidence that differences in performance are attributable to their innovation rather than to differences in training data, hyperparameters, or infrastructure that plague comparisons against proprietary baselines. The model's solid but not saturating performance (53.0 OpenLLM average, well below the ~70+ achieved by much larger models) leaves substantial headroom for proposed improvements to demonstrate gains, making it a useful benchmark target rather than a ceiling.

Instruction-tuning data generation via the two-phase curriculum approach. The paper's Phase 2 strategy—injecting synthetic instruction data during the learning rate decay—suggests a practical pipeline for organizations that need domain-specific models. Rather than pretraining on general web data and then fine-tuning on domain data as separate phases (the standard approach), practitioners could adopt JetMoE's integrated curriculum: pretrain on broad data with a high constant learning rate (Phase 1), then switch to a mixture that includes domain-specific synthetic instructions and textbooks during the decay (Phase 2). The paper's documentation of exact data percentages (Table 2) provides a concrete starting point for designing such curricula. The 60 H100 hours for the dSFT/dDPO alignment phase demonstrates that the post-pretraining fine-tuning can be done at minimal additional cost, making the overall pipeline suitable for small organizations that cannot afford separate large-scale pretraining and fine-tuning runs.

When to Prefer This Method

The paper itself does not articulate an explicit tradeoff framework positioning JetMoE against named alternatives—it presents the model as a demonstration of cost-effectiveness rather than a recommended choice over specific competing architectures. The comparisons against Llama2, DeepseekMoE, and Gemma in Table 3 are benchmarks, not decision rules, and the paper does not specify conditions under which one should prefer JetMoE over these alternatives. The model's availability is presented as a contribution to the open-source ecosystem, not as a substitute for existing models in specific use cases. A forced "Prefer JetMoE when / Prefer Llama2 when" matrix would therefore be imposing a framing the paper does not provide.