PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Diffusion Language Models via Adaptation from Autoregressive Models

Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, Hao Peng, Lingpeng Kong

arXiv:2410.17891v3cs.CL

TL;DR

Diffusion language models have lagged autoregressive models in scale and fair evaluation, while training large DLMs from scratch is challenging. The paper adapts pretrained AR models by unifying their objectives and addressing architectural differences, producing DiffuGPT and DiffuLLaMA. Across language modeling, reasoning, and commonsense evaluation, the adapted models outperform earlier DLMs and are competitive with AR counterparts.

  • Problem

    DLMs remain smaller than AR models and lack fair large-scale benchmark comparisons, while training diffusion models from scratch at scale is challenging.

  • Method

    The paper continually pretrains GPT2 and LLaMA2 into diffusion models by unifying AR and diffusion objectives, annealing attention masks, and retaining the shift operation.

  • Results

    127M-7B AR models become DiffuGPT and DiffuLLaMA with less than 200B training tokens, outperforming earlier DLMs and remaining competitive with AR counterparts.

  • Takeaways & Limitations

    The adapted models extend diffusion language modeling to 7B parameters while supporting in-context learning, code generation, and strong infilling capabilities.

Abstract

from arXiv · show

Diffusion Language Models (DLMs) have emerged as a promising new paradigm for text generative modeling, potentially addressing limitations of autoregressive (AR) models. However, current DLMs have been studied at a smaller scale compared to their AR counterparts and lack fair comparison on language modeling benchmarks. Additionally, training diffusion models from scratch at scale remains challenging. Given the prevalence of open-source AR language models, we propose adapting these models to build text diffusion models. We demonstrate connections between AR and diffusion modeling objectives and introduce a simple continual pre-training approach for training diffusion models. Through systematic evaluation on language modeling, reasoning, and commonsense benchmarks, we show that we can convert AR models ranging from 127M to 7B parameters (GPT2 and LLaMA) into diffusion models DiffuGPT and DiffuLLaMA, using less than 200B tokens for training. Our experimental results reveal that these models outperform earlier DLMs and are competitive with their AR counterparts. We release a suite of DLMs (127M-355M-7B) capable of generating fluent text, performing in-context learning, filling in the middle without prompt re-ordering, and following instructions https://github.com/HKUNLP/DiffuLLaMA.

1 INTRODUCTION

The paper addresses the scale and comparability gap between diffusion language models and autoregressive models by adapting pretrained AR models into DLMs. It evaluates these adapted models broadly and reports competitive capabilities across language modeling, reasoning, and infilling.

  • Diffusion language models offer controllable, any-order, and parallel generation, plus intermediate token correction and global planning.
  • Existing DLMs remain smaller and under-trained than AR models, limiting fair comparisons on advanced capabilities such as mathematical reasoning.
  • The proposed adaptation unifies AR and diffusion objectives, anneals causal attention masks, and retains AR shift operations.
  • The evaluation compares AR and diffusion models across language modeling, reasoning, and infilling in zero-shot, few-shot, and fine-tuning settings.
  • 127M-7B GPT2 and LLaMA2 models become DiffuGPT and DiffuLLaMA using less than 200B training tokens; DiffuGPT outperforms GPT2 in most tasks.
  • DiffuLLaMA adapts a 7B AR model and exhibits in-context learning, code generation, strong infilling, and competitive generation speed.Its reported speed is competitive for generating 1024 tokens with 256 diffusion timesteps.
  • The authors release adapted 127M, 355M, and 7B models with adaptation code, fine-tuning scripts, and evaluation toolkits.

2 PRELIMINARY AND NOTATION

This section introduces diffusion models as forward-noising and reverse-denoising latent-variable models, then specializes the process to discrete text diffusion. It also motivates continuous-time sampling to avoid fixed noise-ratio bias.

  • Diffusion models use a forward Markov process that corrupts data into increasingly noisy variables and a reverse process that denoises them toward the original data.
  • The model parameters are learned by minimizing the negative log-likelihood, optimized through an evidence lower bound.
  • Continuous text diffusion applies Gaussian perturbations across time steps, with β_t controlling noise scales and x_T approaching a standard normal distribution.
  • Discrete denoising diffusion uses categorical transitions over one-hot vocabulary vectors, with transition matrices defining token-state changes.
  • Absorbing discrete diffusion leaves tokens unchanged with probability 1 − β_t and otherwise transitions them into a special [MASK] token.
  • The forward process is expressed as α_t x_0 + (1 − α_t)m, with α_T approaching 0 so fully noised data becomes [MASK] with probability 1.
  • Continuous-time sampling lets t span any point in [0, 1], avoiding the fixed noise ratios imposed by discrete time representations.

3 MODEL

The model connects autoregressive and discrete diffusion objectives, then adapts pretrained autoregressive transformers through masking, shifting, and architecture choices. It trains and samples diffusion models with procedures that preserve useful autoregressive structure while enabling iterative denoising.

  • Unifying Language Modeling Objectives: The discrete diffusion objective becomes a reweighted cross-entropy loss in the continuous-time limit.As T →∞, the first two ELBO terms approach zero and a constant, leaving the per-step loss structure.
  • Unifying Language Modeling Objectives: AR language modeling predicts each ground-truth token from preceding tokens, whereas diffusion predicts masked tokens with reweighted cross-entropy.The diffusion weighting emphasizes smaller t, while AR uses unidirectional attention and shifted targets.
  • Unifying Language Modeling Objectives: AR modeling can be viewed as deterministic right-to-left masking followed by left-to-right denoising over N steps.This establishes an objective-level connection between standard autoregressive generation and a special diffusion process.
  • Adaptation: The adaptation recipe combines attention-mask annealing, shift operations, and a time-embedding-free architecture.These choices narrow the differences between pretrained AR models and diffusion language models.
  • Adaptation: Attention-mask annealing progressively expands right-side context from causal masking to full attention.At each training step, the sampled right-side context increases while preserving left-to-right conditionals during adaptation.
  • Adaptation: The shift operation keeps AR output alignment while training the diffusion model to recover original token positions.Logits are shifted, and prediction targets are aligned so inherited representations remain compatible with diffusion training.
  • Sampling: Sampling starts from all mask tokens, reverses the diffusion process, shifts generated sentences back, and prepends a start token before subsequent passes.The number of timesteps controls a computation–quality trade-off, while conditional sampling retains only high-probability tokens.

4 EXPERIMENT

The experiments evaluate adapted diffusion models across language modeling, reasoning, commonsense, infilling, in-context learning, ablations, and decoding efficiency. The results show strong performance among DLMs, with benefits from scaling, discrete diffusion adaptation, and efficient decoding, while some gaps remain against base AR models.

  • Experimental Setup: 30B tokens for DiffuGPT and 65B tokens for DiffuLLaMA continue pre-training GPT2 and LLaMA-2, respectively.DiffuGPT uses FineWeb, while DiffuLLaMA uses a SlimPajama–Starcoder mixture.
  • Benchmark Performance: Adapted diffusion models achieve state-of-the-art results among existing DLMs, with larger models generally performing better.DiffuLLaMA still falls short of LLaMA2, likely because it was trained on a small data subset; the authors suggest more tokens could help.
  • Benchmark Performance: DiffuLLaMA outperforms AR models on some extensive global-reasoning tasks, while diffusion models show particular strength on infilling.The paper attributes the infilling advantage partly to diffusion training naturally supporting reordered context without the AR model’s FIM preparation.
  • Unconditional Generation: DiffuGPT achieves lower perplexity while maintaining high distinct 2-gram diversity, and quality improves with more decoding steps and larger model size.DiffuGPT outperforms SEDD and MD4 particularly at lower step counts, whereas Plaid 1B needs more decoding steps for fluent text.
  • In-Context Learning: Few-shot performance improves over zero-shot performance for DiffuLLaMA, suggesting that the adapted model can learn from in-context examples.The authors especially observe improved adherence to answer formats and hypothesize that some base AR abilities are retained.
  • Ablation Study: 45.4 and 49.7 accuracy are achieved by direct DD-loss finetuning of GPT2 small and medium, while adapted DiffuGPT reaches 50.2 and 61.8.Direct DD finetuning already outperforms GPT2 AR finetuning, but starting from DiffuGPT improves accuracy and convergence.
  • Decoding Efficiency: DiffuLLaMA achieves better inference efficiency than comparison models at T = 256 for sequences of length 1024 or longer using flash-attention 2.Lowering T can accelerate decoding further but may reduce quality.

5 RELATED WORK

Related work explores adapting language models to new domains or abilities, text diffusion models for generation, and non-autoregressive approaches that relax left-to-right decoding.

  • Continue Pre-training: Continual pre-training adapts existing language models to domain-specific settings or abilities such as longer context and code generation.The cited examples motivate transferring knowledge from existing models when pre-training new architectures is expensive.
  • Text Diffusion Models: Text diffusion research extends diffusion modeling to text generation through continuous and discrete formulations, including small-scale pre-training and fine-tuning frameworks.The passage situates the paper among several lines of text-diffusion work and notes prior small-scale experiments.
  • Non-autoregressive Generation: Non-autoregressive models remove the left-to-right generation constraint, enabling capabilities such as planning with future tokens.Diffusion language models are presented as a notable non-autoregressive family, while other approaches seek practical trade-offs.

6 CONCLUSION

The paper presents a recipe for adapting off-the-shelf autoregressive LLMs into diffusion language models by unifying objectives and modifying attention behavior. Experiments show that DiffuGPT and DiffuLLaMA perform strongly across several tasks, while DiffuLLaMA supports in-context demonstrations to some extent.

  • 6 CONCLUSION: The adaptation recipe continues training off-the-shelf autoregressive LLMs as diffusion language models.It combines objective unification with attention mask annealing and a shift operation.
  • 6 CONCLUSION: Attention mask annealing enables bidirectional modeling, while the shift operation preserves training dynamics similar to autoregressive models.These components are used alongside unified autoregressive and diffusion language-modeling objectives.
  • 6 CONCLUSION: DiffuGPT and DiffuLLaMA show better performance than existing diffusion language models across commonsense reasoning, language modeling, math reasoning, and code generation.The conclusion reports this cross-task comparison without specifying individual benchmark values.
  • 6 CONCLUSION: DiffuLLaMA can follow in-context demonstrations to some extent on math problems.The conclusion also identifies instruction tuning and inference-time planning as future directions.

A OBJECTIVE DERIVATIONS

The appendix introduces diffusion as a forward corruption process paired with a learned reverse denoising process, covering continuous perturbations and discrete masking for text.

  • A OBJECTIVE DERIVATIONS: The forward process transforms data x0 into increasingly noisy variables x1:T, while the reverse Markov process denoises xt toward x0.The model parameters are learned by minimizing the negative log-likelihood through a variational lower bound.
  • A OBJECTIVE DERIVATIONS: Continuous text diffusion applies perturbations at each forward step, whereas discrete denoising models use categorical variables aligned with discrete text.The discrete formulation represents tokens with one-hot vectors and a vocabulary-sized categorical distribution.
  • A OBJECTIVE DERIVATIONS: In the discrete process, tokens remain unchanged with probability 1 − βt and otherwise transition to the absorbing mask state.The transition matrix specifies probabilities between token states, with eK designated as the special mask state.
  • A OBJECTIVE DERIVATIONS: The noise schedule is chosen so that αT approaches 0, making the fully corrupted variable xT equal to the mask state with probability 1.This assumption defines the endpoint of the discrete diffusion process used in later sections.

A.2 LOSS DERIVATION

The appendix extends discrete diffusion to continuous time and sequences, derives denoising transitions and loss terms, and describes training and implementation choices for DiffuGPT and DiffuLLaMA.

  • A.2 LOSS DERIVATION: Continuous-time diffusion allows arbitrary time points by dividing [0, 1] into intervals and taking T →∞.The forward process generalizes from q(xt|xt−1) to q(xt|xs) for 0 ≤ s < t ≤ 1.
  • A.2 LOSS DERIVATION: For the discrete masking process, the marginal distribution is q(xt|x0) = αtx0+(1−αt)m, where m is the absorbing mask state.This marginal supports deriving transitions between arbitrary times and the backward transition probability.
  • A.2 LOSS DERIVATION: The denoising model approximates reverse transitions using pθ(xs|xt, fθ(xt)), with fθ(xt) obtained by applying softmax to transformer logits.The resulting KL-divergence objective connects the denoising model to the reverse diffusion process.
  • A.2 LOSS DERIVATION: The KL-divergence term becomes a cross-entropy form, −x⊤0 log fθ(xt), under the stated conditional formulation.The appendix uses an indicator function to unify cases depending on whether xt equals the mask state.
  • A.2 LOSS DERIVATION: For a sequence of length N, the single-token formulation extends to xt = [x1t … xNt], with training sampling t and evaluation using the integral loss.The sequence loss aggregates token-level modeling over the whole input sequence.
  • A.2 LOSS DERIVATION: DiffuGPT uses 30 billion randomly sampled FineWeb tokens, while DiffuLLaMA uses 65 billion tokens mixed 7:3 from SlimPajama and Starcoder.The datasets are selected to maintain continuity with the base models’ pre-training distributions and minimize distributional shift.
  • A.2 LOSS DERIVATION: DiffuLLaMA is implemented with ZeRO-3, CPU offloading, flash-attention 2, and fused cross-entropy to improve memory and compute efficiency.The setup uses a batch size of 60 per GPU with context length 2048 on a GH200 96GB GPU.
  • A.2 LOSS DERIVATION: The adaptation keeps the base tokenizer unchanged and repurposes an existing low-frequency vocabulary item as the mask token.This avoids expanding the vocabulary with an additional mask-token dimension.

B.3 EVALUATION DETAILS

The evaluation measures conditional completion, commonsense reasoning, infilling, and unconditional generation under specified sampling settings. Generation quality depends on diffusion steps and sampling hyperparameters, with lower temperature improving fluency but reducing diversity.

  • Generation tasks: TriviaQA and LAMBADA use oracle-length continuation settings, while DLM diffusion timesteps match the required generation length.TriviaQA evaluates the first 2000 cases and adds 10 tokens beyond oracle length; LAMBADA completes the final word.
  • CommonSense Reasoning tasks: Commonsense reasoning selects among four choices by comparing diffusion losses for each question-choice pair.A lower loss, interpreted as lower perplexity, indicates the preferred choice.
  • Infilling tasks: ROCStories infilling predicts the third sentence from surrounding context, while HumanEval evaluates single-line code infilling.The evaluation uses 1,000 ROCStories cases and 1,033 HumanEval cases.
  • Unconditional Generation: Unconditional generation is evaluated with GPT-2 large perplexity and distinct 2-gram diversity across 64 generated samples.The reported sampling settings use top-k or top-p temperatures depending on model size.
  • Unconditional Generation: Lowering sampling temperature increases fluency but reduces diversity, producing noticeable sentence repetition.Figure 5 compares generation quality across diffusion timesteps and sampling algorithms.
  • Unconditional Generation: Lower diffusion timestep values produce less fluent generated text in DiffuGPT-M examples.The examples compare T values of 1024, 256, and 32.

C.2 ABLATION ON GSM8K-SYMBOLIC

The GSM8K-symbolic ablation compares direct GPT2 adaptation with DiffuGPT initialization under discrete diffusion objectives. DiffuGPT initialization yields higher accuracy, faster convergence, and lower loss.

  • Accuracy comparison: 50.2 and 61.8 accuracy result from finetuning DiffuGPT small and medium models, exceeding direct GPT2 finetuning.Direct GPT2 finetuning reaches 45.4 and 49.7 for the corresponding small and medium models.
  • Convergence: DiffuGPT converges faster and attains lower loss than GPT2 under the same discrete diffusion loss.The comparison supports DiffuGPT as a stronger diffusion starting point for finetuning.
  • Loss-function ablation: The ablation evaluates DD, DD without shifting, and DD without annealing using training loss and accuracy.It also reports DiffuGPT and DiffuLLaMA results for comparison.

C.3 ADVANTAGES OF DLMS

The experiments highlight DLM advantages in self-correction, global planning, and infilling. These capabilities arise from iterative masked-token refinement and remain useful across diffusion model sizes and tasks.

  • Self-correction: DiffuGPT refines intermediate numbers during later sampling steps without a left-to-right constraint.The qualitative analysis identifies a self-correction capability in its sampling trajectories.
  • Global planning: DiffuGPT demonstrates advantages over AR baselines on the counting-down mathematical reasoning task.The counting-down dataset generalizes the game 24 and is described as challenging for many AR models.
  • Infilling: DLMs are naturally suited to infilling because they are trained to handle masked inputs, unlike AR models that require additional FIM training or careful prompting.The LLaMA prompting attempt performs no better than completing the prefix.
  • Infilling: Diffu-CodeLLaMA outperforms prefix-suffix-middle AR finetuning on HumanEval infilling, indicating prompt format affects AR models but not DLMs.The controlled comparison trains both model types on 100M Starcoder tokens.

C.4 CONTINUAL PRE-TRAINING AR MODELS

Continual pre-training does not improve GPT2’s zero-shot performance in the reported comparison, whereas adapted diffusion models show scaling and inference-efficiency trends. Faster decoding through fewer diffusion steps can reduce downstream accuracy.

  • Continual pre-training: GPT2 continual pre-training on the same corpus and settings as DiffuGPT produces no zero-shot improvement.The authors associate this result with a stability gap and limited new knowledge from a similar corpus.
  • Decoding speed: Smaller diffusion timestep values accelerate generation but may slightly reduce accuracy on downstream multiple-choice tasks.The paper presents this as a speed-quality trade-off controlled through T.
Loading 2410.17891v3…