PaperlayerView exact revision on arXiv

Source-linked AI summary

LLaDA-MoE: A Sparse MoE Diffusion Language Model

Fengqi Zhu, Zebin You, Yipeng Xing, Zenan Huang, Lin Liu, Yihong Zhuang, Guoshan Lu, Kangyu Wang, Xudong Wang, Lanning Wei, Hongrui Guo, Jiaqi Hu, Wentao Ye, Tieyuan Chen, Chenchen Li, Chengfu Tang, Haibo Feng, Jun Hu, Jun Zhou, Xiaolu Zhang, Zhenzhong Lan, Junbo Zhao, Da Zheng, Chongxuan Li, Jianguo Li, Ji-Rong Wen

arXiv:2509.24389v1cs.CLcs.AI

TL;DR

Prior masked diffusion models largely relied on dense Transformer backbones, leaving scratch pretraining with sparse MoE unstudied to the authors’ knowledge. This paper introduces LLaDA-MoE, a diffusion language model trained from scratch with sparse MoE; with 1.4B active parameters, it surpasses prior 8B diffusion models and is comparable to Qwen2.5-3B-Instruct across benchmarks.

  • Problem

    Prior MDM work largely relied on dense Transformer backbones, and the authors knew of no MDM pretrained from scratch with sparse MoE.

    Source: §1 · p. 1

  • Method

    LLaDA-MoE combines masked diffusion modeling with a sparse MoE architecture and is trained from scratch.

    Source: §1 · p. 2

  • Results

    With 1.4B active parameters, LLaDA-MoE surpasses prior 8B-parameter diffusion models and is comparable to Qwen2.5-3B-Instruct across benchmarks.

    Source: §5 · p. 7

  • Takeaways & Limitations

    The results support MoE as an effective foundation for efficient MDMs and leave room for further research and improvements.

    Source: §5 · p. 7

    The study is potentially constrained by the current model size; future work plans to scale LLaDA-MoE and address accompanying challenges.

    Source: §5 · p. 7

Abstract

from arXiv · show

We introduce LLaDA-MoE, a large language diffusion model with the Mixture-of-Experts (MoE) architecture, trained from scratch on approximately 20T tokens. LLaDA-MoE achieves competitive performance with significantly reduced computational overhead by maintaining a 7B-parameter capacity while activating only 1.4B parameters during inference. Our empirical evaluation reveals that LLaDA-MoE achieves state-of-the-art performance among diffusion language models with larger parameters, surpassing previous diffusion language models LLaDA, LLaDA 1.5, and Dream across multiple benchmarks. The instruct-tuned model LLaDA-MoE-7B-A1B-Instruct demonstrates capabilities comparable to Qwen2.5-3B-Instruct in knowledge understanding, code generation, mathematical reasoning, agent and alignment tasks, despite using fewer active parameters. Our results show that integrating a sparse MoE architecture into the training objective of masked diffusion language models still brings out MoE's strengths under efficient inference with few active parameters, and opens ample room for further exploration of diffusion language models. LLaDA-MoE models are available at Huggingface.

1 Introduction

Prior masked diffusion language models largely used dense Transformer backbones, leaving sparse MoE pretraining from scratch unexplored to the authors’ knowledge. LLaDA-MoE combines masked diffusion with sparse MoE and reports strong performance with few active parameters.

1 Introduction

  • To the authors’ knowledge, prior MDMs relied largely on dense Transformer backbones, while scratch-pretrained sparse MoE had not been studied for MDMs.MoE had been validated in autoregressive models, where sparse routing activates only a subset of parameters per token.Source: §1 · p. 1
  • With 1.4B active parameters, LLaDA-MoE surpasses prior 8B-parameter diffusion models and matches Qwen2.5-3B-Instruct across diverse tasks.The benchmark figure covers knowledge, reasoning, mathematics, coding, and agent tasks.Source: Figure 1 · p. 2, §1 · p. 2
  • LLaDA-MoE introduces a diffusion language model with a sparse MoE architecture, trained from scratch to keep the active-parameter budget small during inference.The contribution summary reports state-of-the-art diffusion-language-model performance and comparisons with larger diffusion models and Qwen2.5-3B-Instruct.Source: §1 · p. 2

2 Related Work

Related work presents masked diffusion as an iterative alternative to left-to-right autoregressive decoding, then describes MoE routing as a way to activate only selected experts per token. Prior work has scaled and extended diffusion models, while sparse MoE offers a potential route to resource-efficient capacity.

2.1 Diffusion Language Models

  • MDMs iteratively predict masked tokens and progressively unmask sequences, contrasting with autoregressive models that decode left to right.During training, they reconstruct masked tokens from partially observed context while optimizing a variational lower bound on log-likelihood.Source: §2.1 · p. 2
  • Prior diffusion-language-model work includes scaling, multimodal extensions, more efficient inference strategies, and stronger reasoning and coding capabilities.LLaDA was trained from scratch with 8B parameters and achieved performance comparable to LLaMA3 8B across benchmarks.Source: §2.1 · p. 2

2.2 The MoE Architecture

  • MoE routers select the top k experts for each token, activating only that subset while retaining a larger total pool of parameters.Auxiliary load-balancing losses and per-expert capacity limits are used to prevent routing collapse and balance expert usage.Source: §2.2 · p. 3

3 LLaDA-MoE

LLaDA-MoE pairs a sparse MoE architecture with multi-stage pretraining and annealing, followed by supervised fine-tuning. Its training and inference procedures include variable-length training, expert-routing losses, and iterative masked-token generation.

3.1 Overview

  • The architecture uses RMSNorm, SwiGLU, rotary positional embeddings, and QK-layernorm in its multi-head attention blocks.Key architectural parameters are summarized in Table 2.Source: §3.1 · p. 3, Table 2 · p. 5
  • Training uses two 10T-token pretraining stages, two 500B-token annealing stages, and SFT; the second annealing stage uses an 8k context.Pretraining Stage 2 reweights sampling toward mathematics and code; Annealing Stage 1 uses high-quality text.Source: §3.1 · pp. 3–4, Figure 3 · p. 4

3.2 Training and Inference for LLaDA-MoE

  • Training on variable context lengths uses 1% of pretraining steps with lengths sampled from 8 to 4096 tokens, reducing the train–test context mismatch and improving evaluation metrics.The other 99% of steps retain the default 4k context.Source: §3.2 · p. 5
  • Top-k routing weights selected experts’ outputs, while auxiliary load-balancing and Z-loss weights are 0.01 and 0.001, respectively.Both auxiliary losses decrease rapidly early in pretraining and then stabilize at low magnitudes.Source: §3.2 · p. 5
  • SFT masks response tokens while keeping prompts clean; for multi-turn dialogs, training selects a target turn and recovers its masked response tokens.Appended EOS tokens are included in the target segment and loss.Source: §3.2 · pp. 5–6
  • Inference starts from a fully masked sequence and iteratively reduces noise while sampling tokens at masked positions; semi-autoregressive sampling can decode blocks in sequence.Low-confidence remasking may refine outputs during semi-autoregressive sampling.Source: §3.2 · p. 6

4 Experiments

LLaDA-MoE is competitive against larger diffusion-language-model baselines despite activating only 1B parameters, and its instruct version trails Qwen2.5-3B-Instruct by a small margin. The evaluation spans knowledge, reasoning, mathematics, coding, agent, and alignment benchmarks.

  • The benchmark suite spans knowledge, reasoning, mathematics, coding, agent, and alignment tasks.Source: §4 · pp. 6–7
  • Figure 4 tracks Z-Loss and Load-Balancing Loss during the first 1T pretraining tokens.Source: Figure 4 · p. 7
  • For generative benchmarks other than MDMs with published results, evaluation uses semi-autoregressive sampling with generation length 1024 and block length 64.Reported MDM results are taken directly from published numbers.Source: §4 · p. 7
  • Despite activating only 1B parameters, both Base and Instruct outperform LLaDA-8B, Dream-7B, and LLaDA-1.5 on average; the Instruct model trails Qwen2.5-3B-Instruct by a small margin.Source: §4 · p. 7
  • The authors conclude that masked diffusion language modeling works well with sparse MoE, delivering parameter-efficient performance and leaving room for further exploration.Source: §4 · p. 7

5 Conclusion

LLaDA-MoE combines from-scratch diffusion-language-model training with a MoE architecture. It surpasses prior 8B-parameter diffusion language models with 1.4B active parameters, while its instructed version is comparable to Qwen2.5-3B-Instruct; the authors identify model size as a current constraint.

  • With 1.4B active parameters, LLaDA-MoE surpasses prior 8B-parameter diffusion language models.Source: §5 · p. 7
  • LLaDA-MoE is trained from scratch with a MoE architecture for efficient inference.Source: §5 · p. 7
  • After instruction tuning, LLaDA-MoE is comparable to Qwen2.5-3B-Instruct across all benchmarks.Source: §5 · p. 7
  • The authors identify current model size as a constraint and plan to scale LLaDA-MoE while addressing scaling challenges.Source: §5 · p. 7

Related papers

Loading 2509.24389v1…