Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

On the Reasoning Abilities of Masked Diffusion Language Models

Anej Svete, Ashish Sabharwal

TL;DR

The paper addresses the question of what reasoning problems masked diffusion language models (MDMs) can solve and with what efficiency. By formalizing MDMs as idealized planners and predictors within finite-precision, log-width transformers, it connects MDMs to padded looped transformers (PLTs) and chain-of-thought (CoT) paradigms, leveraging classical circuit complexity results. It proves that MDMs are equivalent to PLTs up to a padding factor and that MDMs with polylogarithmic denoising steps realize the class $NC$, while constant-step MDMs remain within $AC^0$, and it shows that MDMs can simulate CoT (with overhead) and CoT can simulate MDMs (with overhead), yielding a systematic view of the trade-offs between parallelism and sequential reasoning. A key separation result shows that CoT with $O( ext{log }N)$ steps cannot match the parallel reasoning power of MDMs with similar complexity, highlighting the contexts in which parallel diffusion can outperform sequential CoT. Overall, the work clarifies when parallelism in MDMs provides genuine gains, informs the design of reasoning-focused architectures, and links MDM expressivity to well-understood computational classes.

Abstract

Masked diffusion models (MDMs) for text offer a compelling alternative to traditional autoregressive language models. Parallel generation makes them efficient, but their computational capabilities and the limitations inherent to their parallelism remain largely unexplored. To this end, we characterize what types of reasoning problems MDMs can provably solve and how efficiently. We do this by connecting MDMs to the well-understood reasoning frameworks of chain of thought (CoT) and padded looped transformers (PLTs) in the finite-precision log-width setting: We show that MDMs and polynomially-padded PLTs are, in fact, equivalent in this setting, and that MDMs can solve all problems that CoT-augmented transformers can. Moreover, we showcase classes of problems (including regular languages) for which MDMs are inherently more efficient than CoT transformers, where parallel generation allows for substantially faster reasoning.

On the Reasoning Abilities of Masked Diffusion Language Models

TL;DR

The paper addresses the question of what reasoning problems masked diffusion language models (MDMs) can solve and with what efficiency. By formalizing MDMs as idealized planners and predictors within finite-precision, log-width transformers, it connects MDMs to padded looped transformers (PLTs) and chain-of-thought (CoT) paradigms, leveraging classical circuit complexity results. It proves that MDMs are equivalent to PLTs up to a padding factor and that MDMs with polylogarithmic denoising steps realize the class , while constant-step MDMs remain within , and it shows that MDMs can simulate CoT (with overhead) and CoT can simulate MDMs (with overhead), yielding a systematic view of the trade-offs between parallelism and sequential reasoning. A key separation result shows that CoT with steps cannot match the parallel reasoning power of MDMs with similar complexity, highlighting the contexts in which parallel diffusion can outperform sequential CoT. Overall, the work clarifies when parallelism in MDMs provides genuine gains, informs the design of reasoning-focused architectures, and links MDM expressivity to well-understood computational classes.

Abstract

Masked diffusion models (MDMs) for text offer a compelling alternative to traditional autoregressive language models. Parallel generation makes them efficient, but their computational capabilities and the limitations inherent to their parallelism remain largely unexplored. To this end, we characterize what types of reasoning problems MDMs can provably solve and how efficiently. We do this by connecting MDMs to the well-understood reasoning frameworks of chain of thought (CoT) and padded looped transformers (PLTs) in the finite-precision log-width setting: We show that MDMs and polynomially-padded PLTs are, in fact, equivalent in this setting, and that MDMs can solve all problems that CoT-augmented transformers can. Moreover, we showcase classes of problems (including regular languages) for which MDMs are inherently more efficient than CoT transformers, where parallel generation allows for substantially faster reasoning.
Paper Structure (44 sections, 36 theorems, 62 equations, 2 figures)

This paper contains 44 sections, 36 theorems, 62 equations, 2 figures.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Transformers
  4. Masked diffusion language models
  5. Our idealization of masked diffusion models
  6. Theoretical results
  7. Notation.
  8. MDMs are equivalent to padded looped transformers
  9. MDMs and CoT can (inefficiently) simulate each other
  10. A separation between MDMs and CoT transformers
  11. Discussion
  12. Strengths and weaknesses of MDMs.
  13. Equivalence to padded looped transformers.
  14. Generalizations.
  15. Discrepancies.
  16. ...and 29 more sections

Key Result

Theorem 2.1

If assumption:perfect-approximationassumption:uniform-unmasking hold for an LM ${{p}}$, ${{p}}$ cannot compute non-${{{{{{\mathtt{AC}}}^{\mathtt{0}}}}}}$ functions.That is, ${{p}}$ can only implement LMs whose next-symbol logits can be computed by ${{{{{{\mathtt{AC}}}^{\mathtt{0}}}}}}$ circuits liu2

Figures (2)

  • Figure 1: A summary of masked diffusion model expressivity. The colored nodes in the top row correspond to transformer-based computational models and the languages recognized by them. The gray nodes in the bottom row correspond to languages in (${{\mathtt{L}}\textnormal{-uniform}}$) classical complexity classes. ${{{\mathcal{X}}}} \hookrightarrow {{{\mathcal{Y}}}}$ indicates the inclusion of ${{{\mathcal{X}}}}$ in ${{{\mathcal{Y}}}}$, and ${{{\mathcal{X}}}} \leftrightarrow {{{\mathcal{Y}}}}$ indicates equality. Dashed arrows represent strict inclusions. Red arrows denote novel results. $\mathtt{Reg}$ refers to all regular languages.
  • Figure 2: Two strategies for solving a mathematical expression. (a) Parallel: Parallel computation of intermediate values (three steps). (b) Sequential: Step-by-step generation (eleven steps).

Theorems & Definitions (77)

  • Theorem 2.1
  • Theorem 3.1: and
  • Corollary 3.1
  • Corollary 3.2
  • Theorem 3.2
  • Corollary 3.3: with constantly many denoising steps
  • Corollary 3.4: with polylogarithmically many denoising steps
  • Theorem 3.3: can simulate transformers
  • Theorem 3.4: transformers can simulate
  • Corollary 3.5
  • ...and 67 more