Source-linked AI summary
Improved Large Language Diffusion Models
Shen Nie, Qiyang Min, Shaoxuan Xu, Zihao Huang, Yuxuan Song, Yong Shan, Yankai Lin, Wayne Xin Zhao, Chongxuan Li, Ji-Rong Wen
TL;DR
Strong autoregressive models dominate language modeling, while earlier bidirectional diffusion models remained behind them. iLLaDA scales fully bidirectional masked diffusion training from scratch and refines post-training, generation, and evaluation. It substantially improves across general, mathematical, and code benchmarks, though reinforcement-learning alignment and larger-scale matched comparisons remain open.
Problem
Earlier large-scale bidirectional diffusion models remained behind strong autoregressive models, leaving substantial room to improve their language-modeling performance.
Method
iLLaDA is an 8B fully bidirectional diffusion model trained from scratch with 12T-token pre-training, 25B-token SFT, variable-length generation, and confidence-based scoring.
Results
iLLaDA substantially improves over LLaDA across general, mathematical, and code benchmarks and remains competitive with strong autoregressive baselines on several evaluations.
Takeaways & Limitations
The results suggest that fully bidirectional diffusion training from scratch can achieve strong language-modeling performance.
Takeaways & Limitations
The study is limited to the 8B scale, lacks reinforcement-learning alignment, and does not provide a fully matched comparison with autoregressive models.
Abstract
from arXiv · showhide
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.
1 Introduction
iLLaDA addresses the gap between early bidirectional diffusion models and strong autoregressive language models by scaling fully bidirectional masked diffusion training and refining its practical recipe. It substantially improves over prior diffusion models and remains competitive with Qwen2.5 7B on several evaluations.
- iLLaDA is an 8B fully bidirectional masked diffusion language model trained from scratch.
- 12T pre-training tokens and a 25B-token instruction corpus support iLLaDA’s larger-scale training recipe.The instruction corpus is used for 12 SFT epochs.
- Variable-length generation improves inference efficiency, while confidence-based scoring improves multiple-choice evaluation.
- iLLaDA obtains the best average performance among the compared bidirectional diffusion models in both base and instruction-tuned evaluations.
- iLLaDA-Base is slightly stronger than Qwen2.5 7B on average, whereas iLLaDA-Instruct still lags behind Qwen2.5 7B Instruct.
2 Approach
The approach retains masked diffusion training while modifying architecture, data processing, fine-tuning, generation, and evaluation for scalable and efficient language modeling. It also uses confidence-based candidate scoring and variable-length diffusion sampling.
- Pre-training: iLLaDA applies a likelihood-based masked diffusion objective by randomly masking tokens and predicting all masked positions.The loss is computed only for masked tokens, unlike fixed-ratio masked language modeling.
- Pre-training: Grouped-query attention reduces cached key/value memory under cache-style implementations, while tied embeddings reduce parameter count.
- Pre-training: Variable-length training randomly splits sequences, packs examples without common-length padding, and uses variable-length attention kernels.
- Supervised fine-tuning: SFT concatenates formatted instruction examples into a continuous corpus, masks prompts and responses, and optimizes the pre-training objective.The corpus contains approximately 25 billion tokens and is used for 12 epochs.
- Inference and evaluation: For multiple-choice evaluation, iLLaDA repeatedly reveals the highest-confidence ground-truth token among remaining masked positions.The resulting confidence score is a task-specific surrogate rather than a likelihood estimate.
- Inference and evaluation: Variable-length generation appends mask blocks, transfers confident predictions to visible tokens, remasks low-confidence positions, and appends new blocks until stopping.
3 Experiments
Experiments show that iLLaDA substantially improves over prior diffusion models across broad benchmark categories and remains competitive with strong autoregressive baselines. Ablations support confidence-based multiple-choice scoring and longer SFT, particularly for reasoning-heavy tasks.
- iLLaDA substantially improves over LLaDA across general, mathematical, and code benchmarks, with especially large gains on BBH, ARC-Challenge, GSM8K, HumanEval, and MBPP.
- Against Qwen2.5 7B, iLLaDA is competitive and achieves the best reported results on MMLU, BBH, ARC-Challenge, and GSM8K among the compared base models.
- After SFT, iLLaDA outperforms LLaDA and Dream on most benchmarks, with especially pronounced improvements on GSM8K, MATH, and HumanEval.
- Confidence-based scoring improves over the likelihood-style baseline by 1.3 on PIQA, 0.6 on ARC-Challenge, and 2.3 on HellaSwag.
- Performance generally improves as SFT epochs increase, supporting long SFT for iLLaDA, especially on reasoning-heavy benchmarks.The study did not train beyond 12 SFT epochs because of compute constraints.
4 Conclusion and Discussion
The paper concludes that scaling and revising the practical recipe enables strong fully bidirectional diffusion training from scratch. It also identifies missing reinforcement-learning alignment, limited 8B scale, and the absence of a fully matched autoregressive comparison as limitations.
- iLLaDA is an 8B fully bidirectional diffusion language model trained from scratch with pre-training scaled to 12T tokens.
- The practical recipe updates model design, learning-rate scheduling, SFT format, multiple-choice scoring, and variable-length generation.
- Across base and instruction-tuned evaluations, these changes substantially improve performance over LLaDA on general, mathematical, and code benchmarks.
- The study is limited to 8B scale and lacks a fully matched comparison with autoregressive models because compute was allocated to 12T-token pre-training.
A Evaluation Details
The evaluation uses benchmark-specific generation configurations, including different maximum lengths and block sizes. The appendix also reports repetitive reasoning loops in difficult Instruct problems and a termination-oriented mitigation.
- iLLaDA-8B-Base uses open-ended generation for BBH, GSM8K, MATH, HumanEval, and MBPP.
- Semi-autoregressive block sampling hurts HumanEval performance, so Base uses a 512-token maximum generation length and block length.
- iLLaDA-8B-Instruct uses benchmark-specific maximum generation lengths and block sizes, ranging from 3/3 for MMLU-Redux to 4096/32 for MMLU-Pro and MATH.
- Some difficult Instruct problems produce repetitive reasoning loops in which the model repeatedly generates phrases and fails to provide a final answer.
- As generation length increases, the model gradually receives a higher probability of emitting </think> to encourage termination and answer production.