Large Language Diffusion Models — part 1/5

Cleaned-text code points [0, 15520). Bibliography omitted here; source.md includes it.

# Large Language Diffusion Models

###### Abstract

The capabilities of large language models (LLMs) are widely regarded as relying on autoregressive models (ARMs). We challenge this notion by introducing LLaDA , a diffusion model trained from scratch under the pre-training and supervised fine-tuning (SFT) paradigm. LLaDA employs a forward data masking process and a reverse generation process, parameterized by a Transformer to predict masked tokens. It provides a principled generative approach for probabilistic inference by optimizing a likelihood lower bound. Across extensive benchmarks on general tasks, math, code, and so on, LLaDA demonstrates strong scalability and performs comparably to our self-constructed ARM baselines. Remarkably, LLaDA 8B is competitive with strong LLMs like LLaMA3 8B in in-context learning and, after SFT, exhibits impressive instruction-following abilities in case studies such as multi-turn dialogue. Moreover, LLaDA addresses the reversal curse, surpassing GPT-4o in a reversal poem completion task. Our findings show the promise of diffusion models for language modeling at scale and challenge the common assumption that core LLM capabilities discussed above inherently depend on ARMs. Project page and codes: https://ml-gsai.github.io/LLaDA-demo/ .

## 1 Introduction

Large language models (LLMs) [ 1 ] fall entirely within the framework of generative modeling. Specifically, LLMs aim to capture the true but unknown language distribution p data ​ ( ⋅ ) p_{\textrm{data}}(\cdot) by optimizing a model distribution p θ ​ ( ⋅ ) p_{\theta}(\cdot) through maximum likelihood estimation, or equivalently KL divergence minimization between the two distributions: max θ 𝔼 p data ​ ( x ) log p θ ( x ) ⇔ min θ KL ( p data ( x ) | | p θ ( x ) ) ⏟ Generative modeling principles . \displaystyle\underbrace{\max_{\theta}\mathbb{E}_{p_{\textrm{data}}(x)}\log p_{\theta}(x)\Leftrightarrow\min_{\theta}\textrm{KL}(p_{\textrm{data}}(x)||p_{\theta}(x))}_{\textrm{Generative modeling principles}}. (1)

The predominant approach relies on the autoregressive modeling (ARM)—commonly referred to as the “next-token prediction” paradigm—to define the model distribution: p θ ​ ( x ) = p θ ​ ( x 1 ) ​ ∏ i = 2 L p θ ​ ( x i ∣ x 1 , … , x i − 1 ) ⏟ Autoregressive formulation , \displaystyle\underbrace{p_{\theta}(x)=p_{\theta}(x^{1})\prod_{i=2}^{L}p_{\theta}(x^{i}\mid x^{1},\dots,x^{i-1})}_{\textrm{Autoregressive formulation}}, (2) where x x is a sequence of length L L , and x i x^{i} is the i i -th token. This paradigm has proven remarkably effective [ 2 , 3 , 4 , 5 ] and has become the foundation of current LLMs. Despite its widespread adoption, a fundamental question remains unanswered: Is the autoregressive paradigm the only path to achieving the core capabilities of LLMs, such as scalability, in-context learning, and instruction-following?

We argue that the answer is not a simple “yes”. The key insight overlooked previously is: It is the generative modeling principles (i.e., Eq. ( 1 )), rather than the autoregressive formulation (i.e., Eq. ( 2 )) itself, that fundamentally underpin the essential properties of LLMs.

In particular, we argue that scalability is primarily a consequence of the interplay between Transformers [ 7 ] , model size, data size, and Fisher consistency 1 1 1 It suggests the ability to recover the true data distribution with infinite data, a sufficiently large network and optimal training. [ 8 ] induced by the generative principles in Eq. ( 1 ), rather than a unique result of the ARMs in Eq. ( 2 ). The success of diffusion transformers [ 9 , 10 ] on visual data [ 11 ] supports this claim. Furthermore, the instruction-following and in-context learning [ 4 ] capabilities appear to be intrinsic properties of all conditional generative models on structurally consistent linguistic tasks, rather than exclusive advantages of ARMs. In addition, while ARMs can be interpreted as a lossless data compressor [ 12 , 13 ] , any sufficiently expressive probabilistic model can achieve similar capabilities [ 14 ] .

However, certain inherent limitations of LLMs can be directly attributed to their autoregressive nature. For instance, the left-to-right generation process restricts their ability to handle reversal reasoning tasks [ 15 ] , highlighting a representative failure in the generalization capabilities of current models.

Motivated by these insights, we introduce LLaDA (Large Language Diffusion with mAsking) to investigate whether the capabilities exhibited by LLMs can emerge from generative modeling principles beyond ARMs, thereby addressing the fundamental question posed earlier. In contrast to traditional ARMs, LLaDA leverages a masked diffusion model (MDM) [ 16 , 17 , 18 , 19 , 20 ] , which incorporates a forward data masking process and trains a mask predictor to approximate its reverse process. This design enables LLaDA to construct a model distribution with bidirectional dependencies and optimize a variational lower bound of its log-likelihood, offering a principled and previously unexplored perspective on the core capabilities of LLMs discussed above.

We adopt the standard pipeline of data preparation, pre-training, supervised fine-tuning (SFT), and evaluation, scaling LLaDA to an unprecedented language diffusion of size 8B. In particular, LLaDA 8B was pre-trained from scratch on 2.3 trillion tokens using 0.13 million H800 GPU hours, followed by SFT on 4.5 million pairs. Across diverse tasks, including language understanding, math, code, and Chinese, LLaDA demonstrates the following contributions: • LLaDA scales effectively to a compute budget of 10 23 10^{23} FLOPs, achieving comparable results to ARM baselines trained on the same data across six tasks, e.g., MMLU and GSM8K.

• The pre-trained LLaDA 8B Base surpasses LLaMA2 7B Base [ 21 ] on nearly all 15 standard zero/few-shot learning tasks while performing on par with LLaMA3 8B Base [ 6 ] , showcasing effective in-context learning capability.

• LLaDA significantly enhances the ability to follow instructions after SFT, as demonstrated in case studies such as multi-turn dialogue.

• LLaDA effectively breaks the reversal curse [ 15 ] with consistent performance across forward and reversal tasks. Notably, it outperforms GPT-4o in a reversal poem completion task.

## 2 Approach

In this section, we introduce the probabilistic formulation 2 2 2 Here, we focus on the approach of LLaDA. A rigorous formulation of MDM is provided in Appendix A for interested readers. , along with the pre-training, supervised fine-tuning, and inference procedures for LLaDA, as illustrated in Fig. 2 .

### 2.1 Probabilistic Formulation

Unlike ARMs in Eq. ( 2 ), LLaDA defines a model distribution p θ ​ ( x 0 ) p_{\theta}(x_{0}) through a forward process and a reverse process [ 16 , 17 , 18 , 19 , 20 ] . The forward process gradually masks tokens independently in x 0 x_{0} until the sequence is fully masked at t = 1 t=1 . For t ∈ ( 0 , 1 ) t\in(0,1) , the sequence x t x_{t} is partially masked, with each being masked with probability t t or remaining unmasked with probability 1 − t 1-t . The reverse process recovers the data distribution by iteratively predicting masked tokens as t t moves from 1 1 to 0 0 .

The core of LLaDA is a mask predictor , a parametric model p θ ( ⋅ | x t ) p_{\theta}(\cdot|x_{t}) that takes x t x_{t} as input and predicts all masked tokens (denoted as M) simultaneously. It is trained using a cross-entropy loss computed only on the masked tokens [ 18 , 19 , 20 ] : ℒ ( θ ) ≜ − 𝔼 t , x 0 , x t [ 1 t ∑ i = 1 L 1 [ x t i = M ] log p θ ( x 0 i | x t ) ] , \displaystyle\mathcal{L}(\theta)\triangleq-\mathbb{E}_{t,x_{0},x_{t}}\left[\frac{1}{t}\sum_{i=1}^{L}\textbf{1}[x_{t}^{i}=\textrm{M}]\log p_{\theta}(x_{0}^{i}|x_{t})\right], (3) where x 0 x_{0} is a training sample, t t is a continuous random variable drawn uniformly from [ 0 , 1 ] [0,1] , x t x_{t} is sampled from the forward process and L L is the sequence length. The indicator function 1 ​ [ ⋅ ] \textbf{1}[\cdot] ensures that the loss is computed only for masked tokens.

Once trained, we can simulate a reverse process (see Sec. 2.4 for details) parameterized by the mask predictor and define the model distribution p θ ​ ( x 0 ) p_{\theta}(x_{0}) as the marginal distribution induced at t = 0 t=0 . The loss function in Eq. ( 3 ) has been proven to be an upper bound on the negative log-likelihood of the model distribution, making it a principled objective for generative modeling: − 𝔼 p data ​ ( x 0 ) ​ [ log ⁡ p θ ​ ( x 0 ) ] ≤ ℒ ⁡ ( θ ) . \displaystyle-\mathbb{E}_{p_{\textrm{data}}(x_{0})}\left[\log p_{\theta}(x_{0})\right]\leq\mathcal{L}(\theta). (4)

Notably, LLaDA employs a masking ratio that varies randomly between 0 and 1 while BERT [ 22 ] uses a fixed ratio. The subtle differences have significant implications, especially at scale: as shown in Eq. ( 4 ), LLaDA is a principled generative model with the potential to perform in-context learning and instruction-following naturally, akin to LLMs. Moreover, its generative perspective implies strong scalability with large data and models as discussed in Sec. 1 . In addition, MaskGIT [ 23 ] adopts a heuristic training objective, which misses the 1 t \frac{1}{t} term compared to Eq. ( 3 ), and lacks a theoretical link to maximum likelihood. We emphasize that it is precisely the theoretical foundation of maximum likelihood estimation that motivated us to scale discrete diffusion models for language modeling.

### 2.2 Pre-training

LLaDA employs a Transformer [ 7 ] as the mask predictor, similar to existing LLMs. However, LLaDA does not use a causal mask, as its formulation allows it to see the entire input for predictions.

We trained two variants of LLaDA with different sizes: 1B and 8B. We summarize the model architecture of LLaDA 8B and LLaMA3 8B [ 6 ] here, and details are provided in Appendix B.2 . We have ensured consistency in most hyperparameters while making several necessary modifications. We use vanilla multi-head attention instead of grouped query attention [ 24 ] for simplicity, as LLaDA is incompatible with KV caching, resulting in a different number of key and value heads. Consequently, the attention layer has more parameters, and we reduce the FFN dimension to maintain a comparable model size. Additionally, the vocabulary size differs due to a tokenizer [ 4 ] adapted on our data.

The LLaDA model is pre-trained on a dataset comprising 2.3 trillion (T) tokens, adhering to a data protocol that aligns closely with existing LLMs [ 25 , 26 ] , without the incorporation of any special techniques. The data are derived from online corpora, with low-quality content filtered through manually designed rules and LLM-based approaches. Beyond general text, the dataset encompasses high-quality code, math, and multilingual data. Please refer to Appendix B.1 for more details about datasets. The mixing of data sources and domains is guided by scaled-down ARMs. The pre-training process utilizes a fixed sequence length of 4096 tokens, incurring a total computational cost of 0.13 million H800 GPU hours, similar to ARMs of the same scale and dataset size.

For a training sequence x 0 x_{0} , we randomly sample t ∈ [ 0 , 1 ] t\in[0,1] , mask each token independently with the same probability t t to obtain x t x_{t} (see Fig. 2 (a)) and estimate Eq. ( 3 ) via the Monte Carlo method for stochastic gradient descent training. In addition, following Nie et al. [27] , to enhance the ability of LLaDA to handle variable-length data, we set 1% of the pre-training data to a random length that is uniformly sampled from the range [ 1 , 4096 ] [1,4096] .

We adopted the Warmup-Stable-Decay [ 28 ] learning rate scheduler to monitor the training progress without interrupting continuous training. Specifically, we linearly increased the learning rate from 0 to 4 × 10 − 4 4\times 10^{-4} over the first 2000 iterations and maintained it at 4 × 10 − 4 4\times 10^{-4} . After processing 1.2T tokens, we decayed the learning rate to 1 × 10 − 4 1\times 10^{-4} and held it constant for the next 0.8T tokens to ensure stable training. Finally, we linearly reduced the learning rate from 1 × 10 − 4 1\times 10^{-4} to 1 × 10 − 5 1\times 10^{-5} for the last 0.3T tokens. Furthermore, we utilized the AdamW optimizer [ 29 ] with a weight decay of 0.1, a batch size of 1280, and a local batch size of 4 4 per GPU. The 8B experiment was executed once, without any hyperparameter tuning.

### 2.3 Supervised Fine-Tuning

We enhance the capability of LLaDA to follow instructions by supervised fine-tuning (SFT) with paired data ( p 0 , r 0 ) (p_{0},r_{0}) , where p 0 p_{0} is the prompt and r 0 r_{0} denotes the response. This is the simplest and most basic post-training method for LLMs. Technically, this requires to model the conditional distribution p θ ​ ( r 0 | p 0 ) p_{\theta}(r_{0}|p_{0}) instead of p θ ​ ( x 0 ) p_{\theta}(x_{0}) in pre-training.

The implementation is similar to pre-training. As shown in Fig. 2 (b), we leave the prompt unchanged and mask the tokens in the response independently, as done for x 0 x_{0} . Then, we feed both the prompt and the masked response r t r_{t} to the pre-trained mask predictor to compute the loss for SFT: − 𝔼 t , p 0 , r 0 , r t [ 1 t ∑ i = 1 L ′ 1 [ r t i = M ] log p θ ( r 0 i | p 0 , r t ) ] , \displaystyle-\mathbb{E}_{t,p_{0},r_{0},r_{t}}\left[\frac{1}{t}\sum_{i=1}^{L^{\prime}}\textbf{1}[r_{t}^{i}=\textrm{M}]\log p_{\theta}(r_{0}^{i}|p_{0},r_{t})\right], (5) where L ′ L^{\prime} denotes a dynamic length specified later, and all other notations remain the same as before.

Note that this approach is fully compatible with pre-training. Essentially, the concatenation of p 0 p_{0} and r 0 r_{0} can be treated as clean pre-training data x 0 x_{0} , while the concatenation of p 0 p_{0} and r t r_{t} serves as the masked version x t x_{t} . The process is identical to pre-training, with the only difference being that all masked tokens happen to appear in the r 0 r_{0} portion.

The LLaDA 8B model undergoes SFT on a dataset comprising 4.5 million pairs. Consistent with the pre-training process, both data preparation and training follow the SFT protocols utilized in existing LLMs [ 25 , 26 ] , without introducing any additional techniques to optimize LLaDA’s performance. The dataset spans multiple domains, including code, mathematics, and instruction-following. We append | EOS | |\text{EOS}| tokens to the end of short pairs in each mini-batch to ensure equal lengths across all data. We treat | EOS | |\text{EOS}| as a normal token during training and remove it during sampling, enabling LLaDA to control the response length automatically. Please refer to Appendix B.1 for more details.

We train for 3 epochs on the SFT data using a similar schedule to the pre-training phase. The learning rate is linearly increased from 0 to 2.5 × 10 − 5 2.5\times 10^{-5} over the first 50 iterations and then kept constant. During the final 10 % 10\% of iterations, it is linearly reduced to 2.5 × 10 − 6 2.5\times 10^{-6} . Additionally, we set the weight decay to 0.1 0.1 , the global batch size to 256 256 , and the local batch size to 2 2 per GPU. The SFT experiment was executed once, without any hyperparameter tuning.

### 2.4 Inference

As a generative model, LLaDA can sample new text and evaluate the likelihood of candidate text in a diffusion manner instead of the left-to-right autoregressive fashion .