PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling up Masked Diffusion Models on Text

Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang, Qian Liu, Guangtao Zeng, Min Lin, Chongxuan Li

arXiv:2410.18514v3cs.AIcs.CLcs.LG

TL;DR

MDM scalability and effectiveness across language understanding and generation tasks remain underexplored. This paper establishes an MDM scaling law, trains models up to 1.1B parameters, and introduces unsupervised classifier-free guidance; MDMs achieve competitive task performance while addressing some ARM limitations, though gaps remain in scaling and conditional generation.

  • Problem

    MDM scalability and effectiveness in core language tasks, including conditional generation and language understanding, remain underexplored.

  • Method

    The paper trains MDMs up to 1.1B parameters, establishes a scaling law, and introduces unsupervised classifier-free guidance using unpaired data.

  • Results

    MDMs achieve comparable performance to ARMs in key tasks and address reverse-curse and temporal-quality limitations, including outperforming much larger Llama and GPT-3 models in these aspects.

  • Takeaways & Limitations

    MDMs are a promising alternative to ARMs for language modeling at scale.

  • Takeaways & Limitations

    A gap compared to ARMs persists in scaling laws and conditional generation.

Abstract

from arXiv · show

Masked diffusion models (MDMs) have shown promise in language modeling, yet their scalability and effectiveness in core language tasks, such as text generation and language understanding, remain underexplored. This paper establishes the first scaling law for MDMs, demonstrating a scaling rate comparable to autoregressive models (ARMs) and a relatively small compute gap. Motivated by their scalability, we train a family of MDMs with up to 1.1 billion (B) parameters to systematically evaluate their performance against ARMs of comparable or larger sizes. Fully leveraging the probabilistic formulation of MDMs, we propose a simple yet effective unsupervised classifier-free guidance that effectively exploits large-scale unpaired data, boosting performance for conditional inference. In language understanding, the 1.1B MDM outperforms the 1.1B TinyLlama model trained on the same data across four of eight zero-shot benchmarks. Notably, it achieves competitive math reasoning ability with the 7B Llama-2 model on the GSM8K dataset. In text generation, MDMs with 16 times more pre-training time offer a flexible trade-off against ARMs with the accelerated sampling technique KV-Cache: MDMs match ARMs in performance while being 1.4 times faster during sampling. Moreover, MDMs address challenging tasks for ARMs by effectively handling bidirectional reasoning and adapting to temporal shifts in data. Notably, a 1.1B MDM breaks the reverse curse encountered by much larger ARMs with significantly more data and computation, such as 13B Llama-2 and 175B GPT-3. Our code is available at https://github.com/ML-GSAI/SMDM.

1 INTRODUCTION

MDMs are presented as an alternative to ARMs because their bidirectional probabilistic framework may address ARM limitations, while their scalability and effectiveness across core language tasks remain underexplored. The paper studies MDMs up to 1.1B parameters and introduces unsupervised classifier-free guidance for conditional inference.

  • ARMs struggle with bidirectional reasoning and temporal shifts despite their strong sequential scaling properties.
  • MDMs fill masked positions across sequences, enabling flexible bidirectional context modeling.
  • The paper trains MDMs up to 1.1B parameters and establishes the first MDM scaling law.
  • Unsupervised classifier-free guidance uses unpaired data to enhance conditional inference without relying on paired data.
  • 16 times more computation is the MDM-to-ARM training gap, while a 1.1B MDM matches same-sized ARM performance at 1.4 times faster sampling.
  • A 1.1B MDM breaks the reverse curse encountered by much larger 13B Llama-2 and 175B GPT-3 models.

2 MASKED DIFFUSION MODELS ON TEXT

MDMs add noise by randomly masking tokens and learn a reverse process that iteratively recovers masked values. Their formulation removes timestep input, supports conditional clean-data prediction, and enables unsupervised classifier-free guidance.

  • Forward process: MDMs use a forward process that gradually adds noise by randomly and independently masking sentence tokens.
  • Forward process: The mask level is controlled by α_t = 1 − t for the best empirical performance reported by prior work.
  • Reverse process: The reverse process starts from a fully masked sequence and iteratively recovers values for masked tokens.
  • Reverse process: MDMs predict clean-data conditional distributions independently of timestep t, unlike other diffusion models.
  • Training objective: The model approximates clean-token probabilities and optimizes an upper bound on negative log-likelihood.
  • Design rationale: Removing timestep input and minimally modifying Transformers enables scaling and unsupervised classifier-free guidance without paired data.

3 SCALING LAWS FOR MASKED DIFFUSION MODELS

The paper develops a scaling-law analysis for MDMs using IsoFLOP experiments and compares their scaling behavior with ARMs. MDM validation loss follows a similar power-law rate, although MDMs require more computation and use smaller optimal models.

  • Motivation: Prior work lacked a detailed scaling-law exploration for MDMs and a fair comparison with ARMs.
  • IsoFLOP analysis: The study scales compute budgets from 6×10^18 to 10^20 FLOPs while varying parameters N and dataset size D under C = 6ND.
  • Scaling-law formulation: Optimal validation loss is fitted with the power law L = e^β* C^α*.
  • Results: MDM validation loss decreases with compute according to a power law at a rate similar to ARMs.
  • Results: 16 times more computational resources are required by MDMs than ARMs for comparable validation losses.
  • Results: MDMs’ optimal model size is approximately half that of ARMs across compute budgets, indicating similar parameter-capacity utilization.
  • Discussion: The 16-times compute gap is smaller than the 64-times gap reported for continuous diffusion models and may narrow through optimization.

4 UNSUPERVISED CLASSIFIER-FREE GUIDANCE

The paper introduces unsupervised classifier-free guidance for MDMs, allowing conditional inference without paired data by using masked conditioning within the model’s pretrained joint distribution. With paired downstream data, conditional fine-tuning further strengthens this approach and outperforms standard classifier-free guidance.

  • Unsupervised CFG: Unsupervised classifier-free guidance extends CFG to conditional inference without requiring paired data.It introduces a mask sequence matching the condition length, translating the unconditional distribution into conditional format without adding information.
  • Unsupervised CFG: MDMs support unsupervised CFG because pretraining captures the joint distribution of sequence segments and thereby learns their conditional distributions.A paired example can be concatenated into one sequence and modeled through unsupervised training.
  • Unsupervised CFG: The dummy-variable conditional distribution and standard CFG’s unconditional distribution share a similar role in the guidance formulation.This correspondence connects the unsupervised formulation to conventional CFG while avoiding paired-data training.
  • Unsupervised CFG: When paired downstream data are available, fine-tuning the MDM conditional distribution further improves unsupervised CFG and outperforms standard CFG trained on paired data.The approach is presented as especially effective for leveraging large-scale unpaired data.

5 LANGUAGE UNDERSTANDING

The paper finds that MDMs are highly competitive with ARMs in language understanding, while addressing evaluation challenges arising from bidirectional context and variable lengths. Its 1.1B MDM performs strongly across zero-shot benchmarks and achieves math reasoning comparable to 7B Llama-2 with far fewer pre-training FLOPs.

  • MDMs are highly competitive with ARMs of similar model sizes and computational budgets in language understanding.
  • Both ARMs and MDMs perform similarly to random guessing on ARC-c, WinoGrande, and MMLU.The paper leaves possible emergent abilities at larger scales for future work.
  • 1.1% variable-length training data is sufficient to activate handling of variable-length inputs.The authors use variable sequence lengths during training to mitigate the discrepancy between fixed training and variable testing context lengths.
  • The chain rule yields higher accuracy on Obqa and PIQA, whereas Monte Carlo estimation performs better on ARC-e, Hellaswag, RACE, and SIQA.Both methods produce identical results on BoolQ because its answers contain only one token.
  • Unsupervised classifier-free guidance significantly enhances MDM performance across all eight evaluated benchmarks without paired data.The method uses a rescaled conditional distribution because standard classifier-free guidance cannot be applied when paired data is unavailable.
  • The 1.1B MDM outperforms same-data, same-size TinyLlama on four of eight tasks and surpasses 1.5B GPT-2 on six of eight benchmarks.The evaluated benchmarks cover commonsense reasoning and reading comprehension in zero-shot settings, with GSM8K evaluated separately for mathematical reasoning.
  • The 1.1B MDM achieves GSM8K accuracy comparable to 7B Llama-2 while requiring less than 5% of its pre-training FLOPs.Both models are fine-tuned for GSM8K, whereas the other benchmarks are assessed zero-shot.

6 CONDITIONAL LANGUAGE GENERATION

The paper evaluates MDMs on conditional language generation and finds that unsupervised CFG enables a flexible quality–efficiency trade-off against ARMs using KV-cache.

  • Conditional generation setup: A 1.1B MDM provides a more flexible and effective quality-efficiency trade-off during conditional inference than a same-sized ARM with KV-cache.The study evaluates practical conditional generation rather than only unconditional generation and measures running time.
  • Unsupervised CFG: Unsupervised CFG leverages large-scale unpaired data, and conditional fine-tuning further improves it beyond standard CFG trained on paired data.The method combines an unsupervised guidance setup with optional fine-tuning of the conditional distribution.
  • Quality-efficiency trade-off: A 1.1B MDM matches a similarly sized ARM in generation performance while sampling 1.4 times faster.The comparison uses 16 times more MDM pre-training time, while the ARM uses KV-cache; experiments run on one NVIDIA A100-40GB GPU.
  • Quality-efficiency trade-off: Increasing MDM sampling steps can surpass ARM generation quality at the cost of being 1.4 times slower.This exposes an adjustable speed-quality trade-off rather than a single fixed operating point.

7 CHALLENGING TASKS FOR ARMS

The paper tests MDMs on the reverse curse and temporal distribution shifts, reporting advantages over ARMs in bidirectional reasoning and newer data.

  • Breaking the reverse curse: MDMs substantially outperform ARMs and T5 on reverse queries while using fewer parameters, computation, and training data.The MDM uses 10% of the parameters, 1% of the computation, and 10% of the data used by 13B Llama-2.
  • Breaking the reverse curse: MDMs perform similarly to ARMs on same-direction queries while achieving higher scores on reverse queries.The authors attribute this capability to an objective modeling all conditional distributions within the data.
  • Temporal quality degradation: On newer 2024 FineWeb data, MDMs outperform ARMs despite having slightly higher perplexity on the SlimPajama validation set.Both 220M-parameter models were trained on SlimPajama and tested on temporally later data.
  • Temporal quality degradation: The temporal-shift advantage is hypothesized to arise from modeling all conditional distributions rather than only unidirectional dependencies.The exact mechanism remains unclear, so the proposed explanation is explicitly tentative.

8 CONCLUSION

The paper presents MDMs as scalable alternatives to ARMs, while acknowledging remaining gaps and the need to study larger models.

  • Conclusion: MDMs show strong scalability and comparable performance to ARMs in key language-modeling tasks through scaling analysis and unsupervised classifier-free guidance.The conclusion frames these results as evidence for MDMs’ potential at scale.
  • Conclusion: MDMs address reverse curse and temporal quality degradation, outperforming much larger Llama and GPT-3 models on these aspects.The conclusion identifies these as major limitations of ARMs that MDMs can effectively address.
  • Limitations and future work: Scaling laws and conditional generation retain gaps relative to ARMs, motivating future research toward more efficient designs and larger MDMs.The authors identify scaling MDMs to sizes matching advanced ARMs as an important future direction.

ETHICS STATEMENT

The ethics statement describes the benefits and misuse risks of language models and commits to transparency, safety, and fairness research.

  • Potential benefits: Language models could enhance communication, automate tasks, and improve access to information across languages.These benefits are presented as the motivation for improving language models.
  • Risks: Misuse could enable language models to generate false information.The statement identifies false-information generation as a potential risk.
  • Risks: Training on biased datasets could cause generated text to perpetuate those biases.The risk is tied specifically to bias in the training data.
  • Mitigation: The authors commit to transparency and continued research on language-model safety and fairness.The commitment is framed as a mitigation approach for the identified risks.

B.4 ADDITIONAL EXPERIMENTAL DETAILS OF CONDITIONAL GENERATION

This section details conditional-generation setups and contrasts standard CFG with unsupervised CFG. It also reports methods used to address train-test discrepancy and generation evaluation.

  • Setup: The conditional-generation evaluation uses identical optimizer settings for ARMs and MDMs during supervised fine-tuning, with AdamW and cosine learning-rate decay.The maximum learning rate is 2 × 10^-4, the minimum is 2 × 10^-5, and linear warm-up lasts 200 steps.
  • Classifier-free guidance: Unsupervised CFG leverages large-scale pre-trained data for a strong unconditional model and fine-tunes only the conditional distribution on paired data.This differs from standard CFG, which fine-tunes both conditional and unconditional distributions on paired data.
  • Evaluation: MT-Bench scores use first-turn ShareGPT conversations without annealing sampling, and are obtained through the gpt-4o-2024-05-13 API.Both ARMs and MDMs are fine-tuned on first-turn conversations.
  • Train-test discrepancy: Variable-length training and padding mask tokens significantly improve MDM performance on language understanding tasks.The comparison uses chain-rule conditional likelihood without unsupervised CFG.
  • Classifier-free guidance: Standard CFG randomly replaces prompts with masked tokens with probability 10% during fine-tuning to train conditional and unconditional distributions.The replacement represents the unconditional distribution.
  • Classifier-free guidance: Both standard CFG and unsupervised CFG use the rescaled conditional distribution during inference.The rescaled distribution is defined in Eq. (7).

C.2 ADDITIONAL RESULTS OF LANGUAGE UNDERSTANDING

The additional language-understanding results examine discrepancy correction, likelihood evaluation, scaling, conditional guidance, and reverse-curse behavior. Across these analyses, performance depends on the evaluation method and data direction, while unsupervised CFG improves several settings.

  • Train-test discrepancy: Both variable-length training and padding mask tokens significantly improve MDM performance on language understanding tasks.This ablation addresses train-test discrepancy.
  • Likelihood evaluation: The optimal likelihood evaluation method differs across tasks, with chain rule favored for step-by-step reasoning and Monte Carlo estimation for contextual understanding.Chain rule performs better on OBQA and PIQA, whereas Monte Carlo estimation performs better on ARC-e, HellaSwag, RACE, and SIQA; BoolQ is identical.
  • Scaling behavior: MDM performance on language-understanding tasks scales with validation loss, consistently with observations for ARMs.The analysis does not apply train-test discrepancy fixes or unsupervised CFG.
  • Unsupervised CFG: Unsupervised CFG outperforms sampling without CFG using half the sampling steps under equal sampling computation.Additional results also report improved MDM performance on reverse queries with unsupervised CFG.
  • Reverse curse: In NameToDescription reverse-task evaluation, T5 has higher BLEU but lower exact-match accuracy than MDM.T5 responses are similar to the ground truth but differ slightly in a few words.

D EVALUATION METRICS

This section defines the benchmarks used to evaluate language understanding and reasoning. The tasks cover science, commonsense, physical and social interaction, reading comprehension, single-word prediction, and multi-step mathematics.

  • Science reasoning: ARC-Easy evaluates reasoning through elementary-level science questions.It is a subset of the AI2 Reasoning Challenge.
  • Question answering: BoolQ evaluates yes-or-no question answering based on a provided passage.The task measures whether models can answer from passage context.
  • Commonsense reasoning: HellaSwag evaluates commonsense reasoning by asking models to complete a sentence with one of four options.
  • Commonsense reasoning: OpenBookQA assesses subject understanding through multi-step reasoning and integration of additional commonsense knowledge.The dataset is modeled after open-book exams.
  • Physical reasoning: PIQA evaluates physical reasoning by requiring the best solution to problems involving everyday physical scenarios.PIQA stands for Physical Interaction Question Answering.
  • Social reasoning: SIQA evaluates commonsense reasoning about social interactions and the motivations behind human behavior.It presents scenarios requiring reasoning about social interactions.
  • Reading comprehension: RACE evaluates high-school reading comprehension by requiring models to understand and interpret text.
  • Mathematical reasoning: GSM8K evaluates question answering involving multi-step mathematical reasoning through grade-school math word problems.It is a high-quality dataset created for this purpose.
Loading 2410.18514v3…