PaperlayerView exact revision on arXiv

Source-linked AI summary

LLaDA MoE v2: Scaling Mixture-of-Experts Diffusion Language Models

Fengqi Zhu, Shaoxuan Xu, Jingyang Ou, Zebin You, Yipeng Xing, Huabin Liu, Xiaolu Zhang, Jun Zhou, Zhenzhong Lan, Yankai Lin, Wayne Xin Zhao, Jianguo Li, Chongxuan Li, Ji-Rong Wen

arXiv:2608.03457v1cs.AI

TL;DR

MoE diffusion language models lack established scaling evidence, despite their distinct denoising objective and conditioning. This paper characterizes optimization, compute allocation, and expert architecture across scales, then uses the findings to train LLaDA MoE v2, which approaches Qwen3 and outperforms SDAR Chat on seven of eight reasoning and coding benchmarks after SFT alone.

  • Problem

    The scaling behavior and compute-allocation principles of MoE diffusion language models remain uncharacterized, limiting evidence for transferring autoregressive design practices.

  • Method

    The paper empirically characterizes optimization, model–data allocation, and expert-architecture scaling across compute scales, using the findings to guide LLaDA MoE v2 training.

  • Results

    After SFT alone, LLaDA MoE v2 outperforms SDAR Chat on seven of eight reasoning and coding benchmarks and approaches Qwen3 on several benchmarks with 65% as many pretraining tokens.

  • Takeaways & Limitations

    The results provide practical scaling laws and design principles for calibrating MoE diffusion language models rather than directly adopting autoregressive trends.

Abstract

from arXiv · show

Diffusion language models (dLLMs) offer an alternative to autoregressive (AR) language modeling, yet the scaling behavior of Mixture-of-Experts (MoE) dLLMs remains poorly understood. We systematically characterize how optimization hyperparameters, compute allocation, and architecture scale for MoE dLLMs, identifying quantitative differences from scaling trends previously reported for AR models. Specifically, for optimization, the optimal nominal batch size grows faster, while the optimal learning rate decays more rapidly with compute. For model--data allocation, IsoFLOP analysis reveals a slight data-side tilt: the optimal token budget grows faster than activated model-side computation. For MoE architecture, larger scales increasingly favor larger expert pools at fixed activated capacity, while moderate expert granularity remains consistently effective and the preferred fraction of activated capacity assigned to shared experts remains stable across scales. Guided by these findings, we train LLaDA MoE v2, a 30B-A3B dLLM, from scratch on 23.5T tokens. With approximately 65\% as many pretraining tokens as Qwen3, LLaDA MoE v2 approaches Qwen3 on several knowledge, reasoning, and coding benchmarks. After supervised fine-tuning alone, it outperforms SDAR Chat on seven of eight reasoning and coding benchmarks and remains close to Qwen3 on several tasks. These results establish practical scaling laws and design principles for MoE dLLMs.

1 INTRODUCTION

This work addresses the limited understanding of scaling in MoE diffusion language models by systematically studying optimization, compute allocation, and architecture across compute scales. It validates the resulting guidance with LLaDA MoE v2, which approaches Qwen3 30B-A3B on several benchmarks using approximately 65% as many pretraining tokens.

  • Motivation: MoE dLLM scaling remains poorly understood because existing dLLM studies have focused on dense architectures, while MoE designs largely inherit autoregressive practice.MoE architectures can expand capacity beyond dense models at the same computation per token, motivating dedicated scaling studies.
  • Study scope: The study systematically characterizes how optimal optimization hyperparameters, compute allocation, and MoE architectural decomposition change across compute scales.The analysis covers batch size, learning rate, activated model-side computation versus training tokens, routing sparsity, expert granularity, and shared capacity.
  • Validation: LLaDA MoE v2 is a 30B-A3B model trained from scratch on 23.5T tokens to validate practical scaling guidance for MoE dLLMs.The model serves as the scale-up evidence accompanying the controlled sweeps.
  • Results: Using approximately 65% as many pretraining tokens, LLaDA MoE v2 approaches Qwen3 30B-A3B on several knowledge, reasoning, and coding benchmarks.The comparison is made against Qwen3 30B-A3B while retaining substantially fewer pretraining tokens.

2 PRELIMINARIES

The section introduces masked diffusion language models, Mixture-of-Experts Transformers, and scaling-law tools for analyzing compute-dependent design choices. Masked diffusion models corrupt token sequences by replacing tokens with [MASK], then learn iterative denoising generation.

  • Masked diffusion language models: Masked diffusion language models define a discrete diffusion process over token sequences by independently masking tokens at noise level t.Each clean token xi_0 is replaced by [MASK] with probability t and otherwise remains unchanged, producing xt.
  • Masked diffusion language models: A bidirectional Transformer predicts clean tokens from corrupted sequences, minimizing a denoising objective that upper-bounds negative log-likelihood.At inference, generation begins with a fully masked sequence and proceeds through iterative denoising steps.
  • Mixture-of-Experts Transformers: Mixture-of-Experts Transformers decouple total model capacity from per-token computation using ne routed experts and a lightweight router.The router activates the top-na routed experts for each token and combines their outputs.
  • Scaling-law analysis: Scaling laws estimate how optimization hyperparameters, model–data allocation, and architecture should change as compute budget C increases.Power-law fits relate optimal batch size and learning rate to C, while IsoFLOP analysis fixes C and sweeps model-side computation against training data.

3 SCALING LAWS FOR MOE DLLMS

This section develops scaling laws for MoE dLLMs across optimization, compute allocation, and architecture. The resulting framework finds faster batch-size growth, faster learning-rate decay, a slight data-side allocation tilt, and scale-dependent architectural preferences.

  • Scaling framework: The framework calibrates compute-dependent batch size and learning rate, allocates compute between activated model-side computation and data, then tunes activation ratio, expert granularity, and shared-expert ratio.These stages translate compute budgets into optimization settings, allocation frontiers, and MoE architectural choices.
  • Optimization scaling: High compute favors larger nominal batches and faster learning-rate decay, with dLLM fits differing systematically from DeepSeek LLM’s AR scaling law.The fitted recommendation also remains close to the best observed configuration at 6 × 10^20 FLOPs.
  • Compute allocation: The compute-optimal frontiers are M ∗(C) = 0.5152 · C^0.475 and D∗(C) = 1.9411 · C^0.525, indicating a slight data-side tilt.The data-side exponent exceeds the model-side exponent, so additional compute is relatively better spent on training tokens than activated non-embedding FLOPs per token.
  • MoE architecture: At fixed activated model-side budget, larger compute scales favor lower activation ratios, although the second-lowest ratio slightly wins at the smallest budgets.The benefit of sparser activation becomes more pronounced as compute increases.
  • MoE architecture: Expert granularity G = 8–16 is robust, while S = 33.3% remains optimal across compute scales despite differing from several AR MoE designs.Granularity shows no monotonic scaling trend, whereas the shared-expert loss curves are U-shaped with minima at S = 33.3%.

4 TRAINING LARGE-SCALE MOE DLLMS

LLaDA MoE v2 is a 30B-A3B MoE dLLM trained from scratch on 23.5T tokens using scaling-law-guided architecture and staged training. It achieves strong pretrained and SFT results, outperforming SDAR baselines broadly while remaining close to Qwen3 on several tasks.

  • Model and training: 23.5T tokens train LLaDA MoE v2, a 30B-A3B MoE dLLM designed from the derived scaling laws.The model has 30B total parameters and 3B activated per token.
  • Model and training: 128 routed experts use top-8 routing plus one shared expert, yielding 9.09% activated capacity and 33.3% shared capacity.Each layer uses a shared expert of width 4d_expert.
  • Model and training: 20T tokens cover Stages 1–2, followed by 2T annealing tokens, 500B long-context training tokens, and 1T final long-context annealing tokens.Stage 4 expands context from 4K to 32K and raises the RoPE base from 10,000 to 500,000.
  • Pretrained evaluation: 58.60 average across 15 benchmarks makes LLaDA MoE v2 the strongest evaluated dLLM, exceeding SDAR Sci by 3.78 points and smaller dLLMs by at least 12.44 points.Coding gains over SDAR Sci include +16.46 on HumanEval and +7.98 on BigCodeBench.
  • Pretrained evaluation: At substantially lower compute budgets, the 30B-A3B model performs comparably to or better than LLaDA MoE 7B-A1B across knowledge, math, and code benchmarks.The comparison evaluates models across varying training-token budgets.
  • Supervised fine-tuning: After SFT alone, LLaDA MoE v2 surpasses SDAR Chat on four reasoning and three code benchmarks, while remaining close to Qwen3 on AIME 24/25, MBPP, and LiveCodeBench.It also surpasses Qwen3 on MultiPL-E without the additional RL used by Qwen3.

5 RELATED WORK

Diffusion language models have emerged as an alternative paradigm to autoregressive language modeling, while Mixture-of-Experts architectures provide a way to increase model capacity without proportional computational growth.

  • Diffusion language models have recently emerged as a new paradigm for language modeling.
  • Masked discrete diffusion models are viewed as potential alternatives to autoregressive models.
  • Mixture-of-Experts (MoE) and Scaling Laws: MoE architectures selectively activate subsets of fine-grained expert parameters, increasing capacity without proportional computation.

6 CONCLUSION · A SCALING LAWS FOR MOE DLLMS

The paper characterizes scaling laws for optimization, model–data allocation, and expert architecture in MoE dLLMs, showing that autoregressive trends require dLLM-specific calibration. These principles guide LLaDA MoE v2 training, which approaches Qwen3 with fewer pretraining tokens and yields an instruct model that outperforms SDAR Chat on seven benchmarks after SFT alone.

  • 6 CONCLUSION: The study characterizes scaling for optimization hyperparameters, model–data allocation, and expert architecture in MoE dLLMs.It identifies practical design principles for MoE dLLM development.
  • A SCALING LAWS FOR MOE DLLMS: Autoregressive scaling trends provide useful priors but require dLLM-specific calibration.The conclusion emphasizes that AR findings cannot be transferred without adjustment.
  • 6 CONCLUSION: The resulting design principles guide from-scratch training of LLaDA MoE v2 30B-A3B.The model is presented as the practical application of the paper’s scaling analysis.
  • 6 CONCLUSION: LLaDA MoE v2 approaches Qwen3 on several benchmarks with fewer pretraining tokens.The comparison concerns benchmark performance and pretraining-token efficiency.
  • 6 CONCLUSION: After SFT alone, without RL, LLaDA MoE v2 produces an instruct model that outperforms SDAR Chat on seven benchmarks.The conclusion attributes this result to supervised fine-tuning without reinforcement learning.
  • A SCALING LAWS FOR MOE DLLMS: Together, the findings establish practical scaling guidance for MoE dLLMs and motivate LLaDA MoE v2 as its trained demonstration.This point synthesizes the paper’s scaling characterization and model-training outcome.

A.1 COMPUTE ACCOUNTING

The paper defines separate total-parameter, activated-parameter, and activated-computation measures for MoE dLLMs. Its compute accounting includes routing, selected routed experts, shared experts, attention, and training over nominal masked and visible tokens.

  • Parameter and compute definitions: The analysis distinguishes total parameters, activated parameters, and activated computation, with each MoE layer selecting n_a of n_e routed experts plus one shared expert.The architecture uses n_layer Transformer layers, hidden size d_model, sequence length s, expert width d_expert, and shared-expert width d_share.
  • Parameter and compute definitions: Ignoring biases and normalization parameters, total parameters use all n_e routed experts, whereas activated parameters use only the n_a selected experts.Both expressions also include attention projections, the router, and the shared expert.
  • Parameter and compute definitions: Setting d_share = n_s d_expert converts the routed-plus-shared expert terms into equivalent-unit forms based on n_e + n_s for total parameters and n_a + n_s for activated parameters.The equivalent-unit forms are 3d_model d_expert(n_e + n_s) and 3d_model d_expert(n_a + n_s), respectively.
  • FLOPs accounting: One multiply–accumulate counts as two FLOPs, and activated non-embedding FLOPs per token combine attention, routing, selected routed experts, and the shared expert.The backward pass is approximated as twice the forward pass; embedding, LM-head, normalization, nonlinearity, and attention-softmax operations are omitted.
  • Training-token accounting: The scaling experiments define training compute from D nominal tokens, including both masked and visible positions, whose Transformer computation is counted equally.Changing the sampled corruption level changes supervised prediction targets but not the accounted FLOPs.

A.2 MOE IMPLEMENTATION · A.3 HYPERPARAMETER SCALING

The MoE implementation combines GQA with routed and shared SwiGLU experts, using top-score routing, pathway balancing, and fixed auxiliary losses. Hyperparameter scaling evaluates batch size and learning rate across model scales and compute budgets, deriving laws from near-optimal training-loss configurations.

  • A.2 MOE IMPLEMENTATION: Each layer pairs grouped-query attention with a SwiGLU MoE feed-forward block containing routed experts and one shared expert.Architecture experiments may omit the shared expert.
  • A.2 MOE IMPLEMENTATION: Routing is performed independently per token and layer by a linear router that selects the na largest routing scores.Selected routing scores are renormalized before combining expert outputs.
  • A.2 MOE IMPLEMENTATION: The routed and shared pathways are combined as Eshare(h) + λEroute(h), while routed-only configurations omit the shared term and scaling factor.For ns > 0, λ is estimated at initialization by matching expected shared and routed output norms using a gate-scaling heuristic.
  • A.2 MOE IMPLEMENTATION: The router is trained jointly with the model using denoising, load-balancing, and router z-loss terms with fixed coefficients αaux = 0.01 and αz = 0.001.Both auxiliary terms are added directly to the denoising loss and held fixed across scaling and architecture sweeps.
  • A.3 HYPERPARAMETER SCALING: Hyperparameter scaling spans 158M, 1B, and 3.6B models under compute budgets from 10^18 to 3×10^20 FLOPs.All runs use the same pretraining data, sequence length 4096, AdamW, and a 2,000-step linear warmup.
  • A.3 HYPERPARAMETER SCALING: For each model scale, the experiments jointly search global nominal token batch size B and peak learning rate η while fixing architecture and training-token budget.The optimizer uses (β1, β2) = (0.9, 0.95) and weight decay 0.1.
  • A.3 HYPERPARAMETER SCALING: Scaling laws use average loss over the final 0.5% of allocated training FLOPs and include configurations within 0.25% of the minimum loss.These near-optimal configurations enter log–log linear regressions of batch size and learning rate against compute C.

A.4 COMPUTE ALLOCATION SCALING

The study uses IsoFLOP sweeps across compute budgets from 10^17 to 10^20 FLOPs to identify optimal model–data allocations. It evaluates MoE models with varying activated computation, assigns tokens as D = C/M, and fits model-side optima in log–log space.

  • Experimental setup: IsoFLOP sweeps span compute budgets from 10^17 to 10^20 FLOPs under fixed data, denoising objective, and optimizer settings.At each budget, nominal token batch size and peak learning rate follow fitted hyperparameter scaling laws.
  • Experimental setup: Warmup is adapted to each allocation using T_warm = max(0.01T, 100) optimizer steps.T denotes the total optimizer steps implied by the model–data allocation.
  • Allocation analysis: At each compute budget, evaluated MoE models vary in activated non-embedding FLOPs per token M, with training-token budget D = C/M.The allocation with the lowest loss is selected as the empirical optimal point.
  • Allocation analysis: Selected model-side optima are fitted by linear regression in log–log space to derive the corresponding data-side frontier.The procedure uses the lowest-loss allocation point at each compute budget.

A.5 MOE ARCHITECTURE SCALING · B TRAINING LARGE-SCALE MOE DLLMS

The study evaluates MoE architecture choices across five fixed compute budgets using controlled sweeps that preserve activated model-side computation. It selects the lowest-loss candidate for each architectural dimension based on final-training loss.

  • A.5 MOE ARCHITECTURE SCALING: A.5 MOE ARCHITECTURE SCALING: Five reference compute budgets from 6 × 1017 to 6 × 1019 FLOPs anchor the architecture sweeps.The Transformer backbone is fixed while one architectural dimension varies at a time, with the other two held as closely as discrete configurations permit.
  • A.5 MOE ARCHITECTURE SCALING: A.5 MOE ARCHITECTURE SCALING: Candidate architectures preserve the activated model-side budget M ∗(C), ignoring negligible router computation.The activation-ratio sweep keeps activated expert computation per token unchanged while lower A increases total parameter count.
  • A.5 MOE ARCHITECTURE SCALING: A.5 MOE ARCHITECTURE SCALING: Expert-granularity candidates vary dexpert while scaling ne, na, and ns in inverse proportion to preserve A, S, and activated widths.This construction preserves both routed and shared activated widths, nadexpert and nsdexpert.
  • A.5 MOE ARCHITECTURE SCALING: A.5 MOE ARCHITECTURE SCALING: Shared-expert-ratio candidates redistribute fixed activated expert width between shared and routed pathways by varying ns and na.The resulting activation-ratio drift is negligible and does not affect the controlled comparison.
  • A.5 MOE ARCHITECTURE SCALING: A.5 MOE ARCHITECTURE SCALING: Each run’s loss is the average training loss over the final 0.5% of allocated training FLOPs.This definition follows Appendix A.3.
  • A.5 MOE ARCHITECTURE SCALING: A.5 MOE ARCHITECTURE SCALING: At every reference budget, the candidate with the lowest loss is selected independently for each architectural dimension.Selection is performed along each architecture sweep dimension.

B.1 MODEL ARCHITECTURE · B.2 PRETRAINING · B.3 SUPERVISED FINE-TUNING

LLaDA MoE v2 is a 30B-A3B diffusion language model built with a 32-layer MoE Transformer, pretrained from scratch on 23.5T nominal tokens and then fine-tuned on 7M instruction–response examples. Its training uses staged data curation and response-only denoising during supervised fine-tuning, with no reinforcement learning stage.

  • B.1 MODEL ARCHITECTURE: The Transformer has 32 layers, hidden size 3072, GQA with 32 query heads and 4 key-value heads, a 157,184-token vocabulary, and MoE feed-forward blocks throughout.
  • B.1 MODEL ARCHITECTURE: The model uses 128 routed experts, activates 8 per token, and assigns 4 expert-width units to one shared expert, yielding A = 9.09% and S = 33.3%.With G = 8, dexpert = dmodel/4 and dshare = dmodel.
  • B.2 PRETRAINING: LLaDA MoE v2 is pretrained from scratch for 23.5T nominal tokens on broadly collected, filtered, deduplicated, high-quality web text using the denoising objective.The pipeline removes boilerplate and malformed or low-quality documents and filters harmful material.
  • B.2 PRETRAINING: Pretraining follows five stages, including separate 10T-token samples in Stages 1 and 2, a curated 1T-token annealing corpus trained for two epochs, and long-form data reaching 32K-token sequences.Stage 2 places slightly more weight on mathematical reasoning and code data.
  • B.2 PRETRAINING: The full pretraining run consumed approximately 460,000 NVIDIA B200 GPU hours while increasing context length from 4K to 32K.
  • B.3 SUPERVISED FINE-TUNING: Supervised fine-tuning runs for three epochs on 7M instruction–response examples, primarily covering single-turn mathematical reasoning and code generation, after data cleaning, deduplication, harmful-content filtering, and conversation formatting.
  • B.3 SUPERVISED FINE-TUNING: During SFT, denoising masks only response tokens while prompts remain uncorrupted conditioning context; MoE load-balancing and router z-losses remain active with αaux = 0.01 and αz = 0.001.
  • B.3 SUPERVISED FINE-TUNING: All parameters are updated with AdamW during SFT, using a global batch size of 512 sequences, a learning rate from 5.0×10−6 to 1.0 × 10−6, and no reinforcement learning stage.The final checkpoint is used as the instruct model.

B.4 EVALUATION · C M ∗(C) 3D∗(C) B η

LLaDA MoE v2 is evaluated across general, mathematical-reasoning, and code-generation benchmarks using task-appropriate likelihood or generation protocols. The accompanying sweeps specify compute-allocation, training, activation-ratio, expert-granularity, and shared-expert-ratio configurations.

  • B.4 EVALUATION: The evaluation covers general tasks, mathematical reasoning, and code generation, including MMLU, GSM8K, MATH, CRUXEval, MBPP, and MultiPL-E.The general-task suite also includes MMLU-Pro, CEval, CMMLU, HellaSwag, and KorBench; the provided passage continues with additional code benchmarks.
  • B.4 EVALUATION: Base-model multiple-choice results use conditional likelihood, while remaining base-model and all SFT results use conditional generation.Official publication results are preferred when available; otherwise, scores are obtained under the study’s evaluation setup.
  • B.4 EVALUATION: For MMLU, MMLU-Pro, CEval, CMMLU, and HellaSwag, accuracy is determined by selecting the candidate answer with the highest conditional likelihood.AR models use left-to-right conditional log-likelihood, while LLaDA MoE v2 follows likelihood protocols from prior diffusion-language-model studies.
  • B.4 EVALUATION: Conditional-generation evaluation executes extracted code against tests and checks extracted mathematical answers with an equivalence checker.Each model generates a completion from the benchmark prompt using its native generation procedure.
  • B.4 EVALUATION: Instruct models generate at most 1,024 tokens, extended to 4,096 for MATH, OlympiadBench, AIME 2024, and AIME 2025.LLaDA MoE v2 uses semi-autoregressive sampling with block size 64 and denoising steps equal to generation length.
  • C M ∗(C) 3D∗(C) B η: Compute-allocation sweeps enumerate candidate architectures and configurations varying FLOPs, model scale, training-token budget, nominal token batch size, and peak learning rate.Table 8 defines candidate architectures, while Table 9 reports the allocation variables for each configuration.
  • C M ∗(C) 3D∗(C) B η: MoE architecture sweeps train candidates for 3D∗(C) nominal tokens, approximately 3C training FLOPs, with B and η evaluated at 3C.The sweeps cover activation ratios, expert granularity, and shared-expert ratios using configurations in Tables 10–13.
Loading 2608.03457v1…