PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning

Jiacheng Ye, Jiahui Gao, Shansan Gong, Lin Zheng, Xin Jiang, Zhenguo Li, Lingpeng Kong

arXiv:2410.14157v3cs.CLcs.LG

TL;DR

Autoregressive language models struggle with complex reasoning and long-term planning, while strategies such as search or backtracking supervision have drawbacks. The paper analyzes subgoal imbalance and introduces MGDM, a discrete diffusion approach that prioritizes subgoals by difficulty. Across complex tasks, MGDM outperforms autoregressive models, with stronger reported Countdown and Sudoku accuracy without search.

  • Problem

    Autoregressive models struggle with complex reasoning and long-term planning, while search can be costly and backtracking supervision can yield verbose inputs and suboptimal performance.

    Source: §1 · p. 1

  • Method

    The paper analyzes subgoal imbalance and proposes MGDM, which uses discrete diffusion to prioritize subgoals by difficulty during learning.

    Source: §1 · p. 2

  • Results

    Across Countdown, Sudoku, and SAT, MGDM outperforms autoregressive models; Countdown accuracy is 91.5% for MGDM versus 45.8% for its autoregressive counterpart without search.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The findings support exploring diffusion-based approaches for sophisticated language understanding and problem-solving tasks.

    Source: §1 · p. 2

Abstract

from arXiv · show

Autoregressive language models, despite their impressive capabilities, struggle with complex reasoning and long-term planning tasks. We introduce discrete diffusion models as a novel solution to these challenges. Through the lens of subgoal imbalance, we demonstrate how diffusion models effectively learn difficult subgoals that elude autoregressive approaches. We propose Multi-Granularity Diffusion Modeling (MGDM), which prioritizes subgoals based on difficulty during learning. On complex tasks like Countdown, Sudoku, and Boolean Satisfiability Problems, MGDM significantly outperforms autoregressive models without using search techniques. For instance, MGDM achieves 91.5\% and 100\% accuracy on Countdown and Sudoku, respectively, compared to 45.8\% and 20.7\% for autoregressive models. Our work highlights the potential of diffusion-based approaches in advancing AI capabilities for sophisticated language understanding and problem-solving tasks. All associated codes are available at \href{https://github.com/HKUNLP/diffusion-vs-ar}{https://github.com/HKUNLP/diffusion-vs-ar}.

1 INTRODUCTION

The paper presents discrete diffusion as an alternative to autoregressive models for complex reasoning and planning, analyzing why difficult subgoals challenge autoregressive learning. It introduces MGDM and reports stronger results than autoregressive counterparts on complex problem-solving tasks.

1 INTRODUCTION

  • Autoregressive models are described as struggling with complex reasoning, long-term planning, and global coherence.The paper connects these difficulties to challenges in robust problem-solving and adaptable cognition.Source: §1 · p. 1
  • Search at inference can incur significant computational costs, while backtracking supervision can produce verbose inputs and suboptimal performance.These are drawbacks the paper attributes to approaches proposed to address autoregressive constraints.Source: §1 · p. 1
  • The paper reports that discrete diffusion models outperform autoregressive counterparts, particularly on tasks requiring complex planning and reasoning.The introduction describes theoretical and empirical evidence through a synthetic planning task.Source: §1 · p. 2
  • Diffusion models learn difficult subgoals through interrelated views that make each view more manageable.The paper frames this as a multi-view learning explanation for the training objective's advantage.Source: §1 · p. 2
  • MGDM prioritizes subgoals by difficulty during learning, aiming for more effective learning and faster convergence.The paper presents it as an extension to discrete diffusion models.Source: §1 · p. 2
  • 91.5% and 100% accuracy on Countdown and Sudoku, respectively, versus 45.8% and 20.7% for autoregressive counterparts, without search.The paper also reports higher accuracy on SAT, particularly as the number of variables and constraints increases.Source: §1 · p. 2

2 BACKGROUND

The background contrasts autoregressive sequence factorization with discrete diffusion's forward noising and learned reverse denoising processes. The diffusion formulation uses categorical distributions and predefined transition matrices.

2.1 AUTO-REGRESSIVE MODELING

  • Autoregressive modeling factorizes sequence joint probabilities into conditional probabilities over preceding tokens.The section introduces x as a sequence drawn from data distribution q(x) and describes optimizing the model with negative log-likelihood.Source: §2.1 · p. 2

2.2 DISCRETE DIFFUSION MODELING

  • Discrete diffusion corrupts original data through a forward noising process and learns a reverse process that denoises latent variables toward the data distribution.The latent sequence is denoted x1:T, and the forward process is defined through successive transitions.Source: §2.2 · p. 2
  • Figure 1 depicts the planning task, while Figure 2 presents method accuracy at 50k training examples and data requirements for solving subgoals at each planning distance.Figure 2's right panel defines solving as accuracy above 90%.Source: Figure 1 · p. 3, Figure 2 · p. 3
  • The diffusion formulation defines forward and backward distributions categorically, using predefined K × K transition matrices and analytically calculable posterior and KL terms.The passage specifies the transition matrix and category count as Qt and K, respectively.Source: §2.2 · p. 3

3 SUBGOAL IMBALANCE AND MULTI-GRANULARITY DIFFUSION MODELS

The section explains how subgoal difficulty can vary with model parameterization, then examines how diffusion learns difficult subgoals and extends it with MGDM to prioritize them during training.

3.1 SUBGOAL IMBALANCE IN AUTOREGRESSIVE AND DIFFUSION MODELING

  • Planning Distance measures how far ahead the model must plan to choose the next node, and increasing it makes that choice harder.The proposition states that subgoals can require substantially more data or be infeasible to learn, depending on the model parameterization.Source: §3.1 · p. 3
  • The diffusion model achieves perfect accuracy across all planning distances, while autoregressive models barely beat random guessing on larger distances and teacherless training produces illegal paths.In this synthetic task, AR and Reverse AR work mainly at PD 0 or 1, where little planning is needed.Source: §3.1 · p. 4
  • AR and Reverse AR learn PD 0 and 1 with 10k examples, but increasingly difficult subgoals require exponentially more data; AR models from 6M to 1.5B parameters still fail to resolve all PDs.Only fine-tuning LLaMA 7B resolves all PD subgoals in the reported parameter-scaling experiment.Source: §3.1 · p. 4

3.2 EFFECTIVE HARD SUBGOAL LEARNING IN DIFFUSION MODELING

  • For PD 3, diffusion’s overall loss is relatively low compared with AR’s, consistent with its stronger performance on difficult subgoals.The figure compares losses for PD 3; the diffusion analysis also examines unweighted loss across timesteps and contexts.Source: §3.2 · p. 5, Figure 3 · p. 5
  • Diffusion treats different noisy versions of the data as interrelated views; less-noisy objectives are easier, and these views make hard subgoals more learnable.When noise masks positions after the hard token, diffusion learns that subgoal more manageably than AR does.Source: §3.2 · p. 5

3.3 MULTI-GRANULARITY DIFFUSION MODELING

  • MGDM extends discrete diffusion by prioritizing difficult views and subgoals during learning to improve learning outcomes and convergence speed.It adds adaptive token-level reweighting: for β > 0, easy-token loss is reduced while harder-token loss is emphasized.Source: §3.3 · pp. 5–6
  • MGDM uses easy-first TopK decoding at inference, which outperforms the random decoding method used by Austin et al. (2021).Source: §3.3 · p. 6

4 EXPERIMENTS

The experiments assess diffusion on complex planning problems; the Countdown subsection describes the task and setup and identifies Table 1 as results across increasing task complexity.

4 EXPERIMENTS

  • The experiments assess the model on three complex problem-solving tasks requiring deliberate planning.The passage refers to the tasks without naming them in the supplied text.Source: §4 · p. 6

4.1 COUNTDOWN

  • Countdown asks solvers to combine given numbers with arithmetic operations to reach a target number.It is described as a generalized version of the game of 24.Source: §4.1 · p. 6
  • The Countdown setup uses 500k generated problems, holds out 10% of targets for out-of-distribution evaluation, and varies input length across 3, 4, and 5 digits.Table 1 reports Countdown results across increasing complexity; the supplied caption does not provide its values.Source: §4.1 · p. 6, Table 1 · p. 6

MGDM (Ours)

MGDM is evaluated on Countdown, Game of 24, Sudoku, and Boolean satisfiability, with analyses of training strategies, decoding speed, and reasoning errors. The reported results favor diffusion over autoregressive models across the problem-solving tasks.

MGDM (Ours)

  • Diffusion approaches outperform autoregressive models across all three Countdown tasks, especially as complexity increases; on Game of 24, 85M MGDM outperforms same-sized GPT-2 and uses 186 times fewer tokens than GPT-4 Tree-of-Thought.Stream-of-Search trajectories can span 60,000 tokens on Countdown 5.Source: MGDM (Ours) · p. 7, Table 2 · p. 7

4.2 SUDOKU

  • A 6M MGDM model perfectly solves the Sudoku test problems, while the 13B LLaMA model achieves only 32.9 performance.Figure 4’s left panel reports Sudoku accuracy.Source: Figure 4 · p. 7, §4.2 · p. 8

4.3 BOOLEAN SATISFIABILITY PROBLEM

  • As SAT variable count increases, accuracy declines for both models, but the performance gap widens in MGDM’s favor.Figure 4’s right panel shows accuracy as Boolean satisfiability difficulty increases.Source: Figure 4 · p. 7, §4.3 · p. 8

4.4 ANALYSIS

  • Token reweighting yields greater improvement than sequence reweighting, easy-first TopK decoding beats random decoding, and multi-granularity loss performs especially well mid-training.Figure 5 tracks evaluation accuracy during training across AR and MGDM reweighting strategies.Source: §4.4 · pp. 8–9, Figure 5 · p. 8, Table 3 · p. 8
  • Figure 6(a) compares accuracy and samples per second as diffusion timesteps vary on Countdown 4 and 5.Countdown 4 and Countdown 5 use the left and right y-axes, respectively.Source: Figure 6 · p. 9
  • On Countdown 4, AR’s calculation errors concentrate in the final equation, with 48.9% there versus 0.2% in the first and 7.2% in the second.The authors call this “The Regretful Compromise”: earlier planning errors leave AR forcing a final equation toward the target.Source: §4.4 · p. 9, Figure 6 · p. 9

5 RELATED WORK

The related work describes autoregressive models and methods developed to address their limitations, then situates discrete diffusion among non-autoregressive approaches. It frames this paper's comparison of diffusion and autoregressive models around subgoal imbalance and complex reasoning and planning.

5.1 AUTOREGRESSIVE MODELING

  • Autoregressive next-token prediction is described as ill-suited to planning, while some next-token prediction problems are computationally hard to approximate.The section also surveys training-stage changes and search-augmented decoding proposed to compensate for autoregressive modeling deficiencies.Source: §5.1 · p. 10
  • Autoregressive-model research includes reverse training, fill-in-the-middle training, future-token prediction, lookahead attention, and search-augmented decoding.The section notes that difficult-to-learn tokens have also been reported in autoregressive pretraining, implying subgoal imbalance in general text corpora.Source: §5.1 · p. 10

5.2 NON-AUTOREGRESSIVE MODELING

  • Traditional non-autoregressive models tend to underperform autoregressive models in generation quality, while text diffusion has approached or outperformed them on several reported measures and tasks.The section presents diffusion as iterative denoising and compares it with autoregressive models through the lens of subgoal imbalance.Source: §5.2 · p. 10

6 CONCLUSION

The paper analyzes autoregressive language-model limitations in deliberate planning and proposes MGDM, reporting stronger performance than autoregressive and previous diffusion models on sophisticated planning tasks.

6 CONCLUSION

  • MGDM performs significantly better than autoregressive and previous diffusion models on various sophisticated planning tasks.The paper describes its analysis of autoregressive limitations in planning across controlled settings and real-world contexts.Source: §6 · p. 10
  • The conclusion calls for reevaluating sequence modeling for modern language models, especially on challenging problem-solving tasks.Source: §6 · p. 10
  • The paper's analysis concerns autoregressive models applied to deliberate planning in controlled settings and real-world contexts.Source: §6 · p. 10

A MORE BACKGROUND AND DERIVATION OF DISCRETE DIFFUSION MODELS

The appendix reviews discrete diffusion and derives its training objective, including the forward corruption and learned reverse process, transition choices, and token-level KL terms.

A.1 BACKGROUND

  • Discrete diffusion developed from binary to categorical data, with later work introducing transition variants and methods for the reverse process.The cited background also notes continuous-time formulations and score-function approaches.Source: §A.1 · p. 18

A.2 DERIVATION SETUP

  • The setup corrupts a one-hot discrete variable through a forward process and learns a backward process that denoises latent variables toward the data distribution.Both transitions are categorical, and the forward transition uses a predefined matrix Q_t.Source: §A.2 · p. 18

A.3 THE MARGINAL AND POSTERIOR

  • The marginal and posterior use the forward process's Markov property, while efficient transition-matrix choices address the cost of cumulative products for large T and category counts.The passage identifies two commonly used forms of Q that allow Q_t to be computed efficiently.Source: §A.3 · p. 18

A.4 TRANSITION MATRIX

  • Absorbing noise sends the corruption process toward a point mass on the absorbing state rather than a uniform distribution.At each step, the process transitions to another token with probability β_t and stays the same with probability 1 − β_t; after t steps, α_t is the probability of not yet transitioning.Source: §A.4 · pp. 18–19

A.5 DERIVATION OF ELBO

  • The derivation introduces a variational upper bound on negative log-likelihood as the objective for fitting the generative model to the data distribution.Source: §A.5 · p. 19

A.6 DERIVATION FOR EQUATION (6)

  • For Equation (6), the derivation expresses the reverse-step KL divergence using an indicator for whether x_t differs from x_0, then extends the loss to sequences with λ_t as the reweighting term.The indicator is 1 when x_t ≠ x_0 and 0 otherwise.Source: §A.6 · p. 20

B ALGORITHMS FOR TRAINING AND INFERENCE

The algorithms describe MGDM training with sampled data, timestep, and corrupted sequence, and sampling by progressively updating masked tokens. Conditional training and inference keep the source portion fixed.

B ALGORITHMS FOR TRAINING AND INFERENCE

  • For conditional training and inference, the source and target are split and the source condition is frozen throughout both procedures.Source: B ALGORITHMS FOR TRAINING AND INFERENCE · p. 20
  • MGDM training samples data, a timestep, and a corrupted sequence, then minimizes a loss using sequence and token reweighting terms.The algorithm takes a custom sequence reweighting term and token reweighting parameters as inputs.Source: B ALGORITHMS FOR TRAINING AND INFERENCE · p. 20
  • MGDM sampling initializes each token as the mask and iteratively selects top-ranked positions for predicted-token updates as timesteps decrease.The algorithm returns the generated sequence after the iterative sampling procedure.Source: B ALGORITHMS FOR TRAINING AND INFERENCE · p. 21

C.1 TASK DETAILS

The appendix describes the datasets and model configurations used for MGDM, autoregressive baselines, and diffusion baselines.

C.1 TASK DETAILS

  • Table 4 summarizes dataset statistics, while Table 10 provides example inputs and outputs for each dataset.Minimal and CD denote the minimal planning task and Countdown, respectively.Source: Table 4 · p. 21, §C.1 · p. 21

C.2 MGDM IMPLEMENTATION DETAILS

  • MGDM uses GPT-2 models trained with task-dependent epochs and sampling steps, a batch size of 1024, and decoding temperature τ = 0.5.Experiments use eight NVIDIA V100-32G GPUs; the authors report checking statistical significance through multiple runs.Source: §C.2 · p. 21

C.3 BASELINE IMPLEMENTATION DETAILS

  • The AR baseline is trained to convergence with task-specific epoch counts, while diffusion baselines use GPT-2 and MGDM training parameters.LLaMA uses LoRA rank 16; GPT-4 numbers are borrowed from Yao et al. (2024).Source: §C.3 · p. 22

D ADDITIONAL EXPERIMENTS

Additional experiments report game-of-24 accuracy and token consumption, and test whether token reweighting improves AR on minimal planning.

D.1 TOKEN CONSUMPTION ON GAME OF 24

  • Table 6 reports detailed game-of-24 accuracy and token consumption.The supplied passages do not state the table's numerical results.Source: §D.1 · p. 22, Table 6 · p. 22

D.2 AR WITH TOKEN REWEIGHTING

  • AR with token reweighting still achieves only around 50% accuracy on subgoals with PD larger than 1, similar to the original AR model.The result is reported for the minimal planning task in Table 7.Source: §D.2 · p. 22, Table 7 · p. 22

D.3 SCALING BOTH DATA AND MODEL SIZE

The appendix reports scaling results for AR and MGDM and provides prediction examples for examining model outputs and errors.

D.3 SCALING BOTH DATA AND MODEL SIZE

  • Scaling both data and model size is effective for both AR and MGDM.Table 8 contains the corresponding results.Source: §D.3 · p. 23, Table 8 · p. 23

D.4 CASE STUDY

  • Table 9's Countdown 4 examples show more AR calculation errors in the last equation following incorrect planning in earlier steps.The table marks planning errors in red and calculation errors in bold.Source: Table 9 · p. 23, §D.4 · p. 23
  • Table 10 provides example inputs and outputs for each dataset.Source: Table 10 · p. 24

Related papers

Loading 2410.14157v3…