Improved Large Language Diffusion Models
AuthorsShen Nie, Qiyang Min, Shaoxuan Xu, Zihao Huang, Yuxuan Song, Yong Shan, Yankai Lin, Wayne Xin Zhao, Chongxuan Li, Ji-Rong Wen
This work shows that large language models can be trained with masked diffusion instead of next-token prediction and still become highly competitive on reasoning, math, and coding tasks.
Key results
iLLaDA is trained as an 8B masked diffusion language model
iLLaDA scales pre-training to 12T tokens
SFT uses a 25B-token instruction corpus
Base model improvement over LLaDA on BBH
Base model improvement over LLaDA on ARC-Challenge
Instruction-tuned model improvement over LLaDA on MATH
What the paper found
Improved Large Language Diffusion Models introduces iLLaDA, an 8B fully bidirectional masked diffusion language model trained from scratch at ByteDance Seed with Renmin University of China. The paper keeps the masked diffusion objective throughout pre-training and supervised fine-tuning, but scales pre-training to 12T tokens, fine-tunes on a 25B-token instruction corpus for 12 epochs, and adds grouped-query attention, tied embeddings, variable-length generation, and confidence-based scoring for multiple-choice evaluation. These changes substantially lift performance over LLaDA: the base model gains 21.6 points on BBH and 14.9 on ARC-Challenge, while the instruction-tuned model gains 14.5 on MATH and 16.5 on HumanEval. On the base suite, iLLaDA reaches 74.8 on MMLU, 81.9 on GSM8K, and 50.0 on HumanEval, and on the instruction suite it reaches 89.0 on GSM8K, 56.7 on MATH, and 65.9 on HumanEval. The ablation shows confidence-based scoring improves multiple-choice accuracy by 1.3 on PIQA, 0.6 on ARC-Challenge, and 2.3 on HellaSwag versus likelihood scoring. Overall, the study argues that fully bidirectional diffusion training from scratch can be competitive with strong autoregressive models such as Qwen2.5 7B on several reasoning benchmarks.
Original abstract
Modern large language models are predominantly trained with autoregressive factorization and causal attention. We present \emph{iLLaDA}, an 8B masked diffusion language model trained from scratch with fully bidirectional attention. iLLaDA keeps the masked diffusion objective throughout pre-training and supervised fine-tuning (SFT), scaling pre-training to 12T tokens and fine-tuning on a 25B-token instruction corpus for 12 epochs. We further use variable-length generation for efficiency and introduce confidence-based scoring for multiple-choice evaluation. Compared with LLaDA, iLLaDA improves broadly across general, mathematical, and code benchmarks; for example, iLLaDA-Base improves by 21.6 points on BBH and 14.9 points on ARC-Challenge, while iLLaDA-Instruct improves by 14.5 points on MATH and 16.5 points on HumanEval. Despite its non-autoregressive training, iLLaDA also remains competitive with Qwen2.5 7B on several benchmarks. These results show that fully bidirectional diffusion training from scratch is a competitive path toward strong language models. Model weights and codes: https://github.com/ML-GSAI/LLaDA.
Read the original paperMore in Large Language Models
Browse all 81 papers →Finetuning with Sampling: SFT Learns Better Than You Think
Aayush Karan, Sitan Chen, Yilun Du
By sampling and reshaping expert data before training, this work argues that supervised finetuning can match RL while generalizing better and forgetting less.
Generalization Dynamics of LM Pre-training
Jiaxin Wen, Zhengxuan Wu, Dawn Song, Lijie Chen
Language models may repeatedly switch between shallow memorization and genuine reasoning during training, and the paper shows how to detect and potentially control these swings.
Rethinking Self-Distillation for Multi-Teacher Capability Merging
Roy Xie, Dan Friedman, Feng Nan, Yukun Huang, Zhichao Xu, Chengjiu Zhang, Jun Xu, Manaal Faruqui, Vivek Rathod, Bhuwan Dhingra
The study finds that expensive multi-teacher on-policy distillation may offer little advantage over carefully tuned, cheaper alternatives such as SFT and weight merging.