Improved Large Language Diffusion Models

Abstract: Modern LLMs are predominantly trained with autoregressive factorization and causal attention. We present \emph{iLLaDA}, an 8B masked diffusion LLM 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 LLMs. Model weights and codes: https://github.com/ML-GSAI/LLaDA.

• The paper introduces iLLaDA, an 8B parameter masked diffusion language model employing bidirectional attention to enhance inference efficiency.
• It utilizes grouped-query attention and tied input/output embeddings to reduce memory footprint and optimize parameter usage.
• Empirical evaluations demonstrate marked gains on benchmarks like BBH, ARC, MATH, and HumanEval compared to autoregressive systems.

Improved Large Language Diffusion Models: A Technical Review

Overview of Methodological Innovations

The manuscript introduces iLLaDA, an 8B parameter masked diffusion LLM employing fully bidirectional attention. Departing from the prevalent autoregressive paradigm, iLLaDA leverages a masked diffusion objective for both pre-training and supervised fine-tuning (SFT). Pre-training is scaled to 12T tokens, and SFT is performed on a 25B-token instruction corpus for 12 epochs. Key architectural modifications include grouped-query attention (GQA) to minimize memory overhead during inference, tied input/output embeddings to reduce parameter count, and sequence length expansion up to 8192. The learning rate schedule is adjusted for large-scale optimization, with cosine decay phases embedded in both pre-training and SFT.

Variable-length generation is adopted for inference efficiency, aligning with recent research on block-wise decoding in masked diffusion models. For evaluation, confidence-based scoring is implemented for multiple-choice tasks, replacing likelihood-based scoring and empirically yielding superior alignment with task objectives.

Diffusion Language Modeling: Objective and Architecture

iLLaDA preserves the masked diffusion formulation previously explored in LLaDA, sampling a masking ratio per sequence and corrupting the data accordingly. The loss is computed solely for masked positions, with the model trained to reconstruct masked tokens given the corrupted context. This approach offers advantages over fixed-ratio masking (e.g., BERT), supporting versatility in masking spans and enabling bidirectional information flow during both training and inference.

The backbone consists of a dense Transformer leveraging RMSNorm, SwiGLU, RoPE, and omitting attention/MLP biases. The switch from multi-head attention to GQA is motivated by practical considerations: GQA reduces the memory footprint of cache states, essential for highly scalable inference in diffusion LLMs with blockwise, semi-autoregressive sampling. The architecture also ties input embeddings and LM-head parameters to optimize parameter efficiency.

Supervised Fine-Tuning and Generation Protocol

SFT in iLLaDA is reformulated: rather than padding responses to uniform length and masking only response segments, iLLaDA applies random masking across entire prompt-response-EOS sequences, sampled from the instruction corpus and batched in variable lengths. This modification ensures that training and inference are closely aligned, enables more robust handling of variable-length generation during sampling, and streamlines blockwise decoding.

The model is fine-tuned for 12 epochs, with extended data reuse yielding significant gains on reasoning benchmarks, as detailed below.

Empirical Evaluation and Numerical Results

A comprehensive evaluation suite covers general language tasks (MMLU, BBH, ARC-Challenge, HellaSwag), mathematics (GSM8K, MATH), and code generation (HumanEval, MBPP). iLLaDA exhibits substantial improvement over baseline bidirectional diffusion models (LLaDA, Dream) in both base and instruction-tuned settings. Key numerical advances include: iLLaDA-Base improving by 21.6 points on BBH and 14.9 points on ARC-Challenge relative to prior diffusion models; iLLaDA-Instruct improving by 14.5 points on MATH and 16.5 points on HumanEval. Notably, iLLaDA-Base achieves parity or outperforms Qwen2.5 7B (autoregressive) on several core benchmarks, demonstrating competitive generalization despite non-autoregressive training.

Ablation on scoring rules supports confidence-based evaluation for multiple-choice tasks, with 1.3–2.3 point improvements on PIQA, ARC-Challenge, and HellaSwag. SFT epoch ablation reveals continuous performance gains on GSM8K, MATH, and MMLU-Pro with extended fine-tuning epochs, evidencing substantial data-reuse potential.

Figure 1: SFT epoch ablation on GSM8K, MATH, and MMLU-Pro, indicating persistent performance improvements across increased fine-tuning duration.

Theoretical and Practical Implications

The results reinforce the assertion that fully bidirectional diffusion models can achieve robust language modeling capabilities, encompassing in-context learning and instruction-following, without reliance on autoregressive factorization. iLLaDA’s architectural and training modifications—most notably GQA, variable-length batching, and confidence-based evaluation—substantially enhance bidirectional masked diffusion scaling, reducing memory requirements and improving learning efficiency.

Data-reuse during SFT emerges as a critical factor, with diffusion models exhibiting resilience to repeated training on smaller corpora, consistent with prior observations in pre-training under data-constrained settings [prabhudesai2025diffusion, ni2025diffusion]. This property may facilitate more efficient adaptation in practical deployment scenarios where instruction or alignment data is limited.

The competitive performance of iLLaDA relative to strong autoregressive baselines (Qwen2.5 7B) signals that diffusion-based pre-training from scratch is viable for high-capacity LLMs, especially in settings requiring bidirectional reasoning, reversal, or multimodal grounding, as indicated by referenced multimodal diffusion research.

Future Directions and Limitations

The authors identify two primary limitations. First, iLLaDA has not yet integrated reinforcement-learning-based alignment (RLHF, VRPO, diffu-GRPO, MDPO, ESPO), which is known to further enhance instruction-following and reasoning performance in autoregressive settings. This omission accounts for the residual gap between iLLaDA-Instruct and autoregressive instruct models. Second, the study is restricted to 8B parameter scale, without full parity in model scaling or reinforcement alignment; larger-scale studies and RL-based refinement of iLLaDA are promising avenues for future research.

Conclusion

iLLaDA advances the practical and theoretical boundaries of diffusion-based language modeling, introducing architectural and training improvements that yield strong general, mathematical, and code modeling performance. Empirical results affirm that fully bidirectional masked diffusion LLMs are competitive with autoregressive counterparts on standard benchmarks. The demonstrated resilience to data reuse, improvements in evaluation via confidence-based scoring, and scalable architecture modifications collectively suggest that diffusion LLMs warrant further exploration for both foundation and specialized LLM deployment.