PaperlayerView exact revision on arXiv

Source-linked AI summary

DINGO: Constrained Inference for Diffusion LLMs

Tarun Suresh, Debangshu Banerjee, Shubham Ugare, Sasa Misailovic, Gagandeep Singh

arXiv:2505.23061v1cs.LGcs.PLcs.SE

TL;DR

Diffusion LLMs predict token blocks in parallel, while existing constrained decoding methods are designed for sequential generation and do not provide the needed guarantees. DINGO uses dynamic programming to enforce regular-expression constraints while selecting the most probable valid output. The paper reports improvements over unconstrained inference on symbolic math and JSON tasks, with up to a 68% improvement on challenging benchmarks.

  • Problem

    Diffusion LLMs lack provable enforcement of user-specified formal constraints, while their parallel token prediction makes existing sequential constrained-decoding methods incompatible.

    Source: §1 · pp. 1–2

  • Method

    DINGO uses dynamic programming to select the maximum-probability valid output under a user-specified regular expression while ensuring constrained outputs are valid prefixes.

    Source: §1 · p. 2, §3 · pp. 3–4

  • Results

    DINGO significantly improves over unconstrained inference on symbolic math and JSON tasks, with up to a 68% improvement reported on challenging benchmarks.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    DINGO provides a dynamic-programming approach for diffusion LLM outputs that adhere to regular-language constraints while preserving the model’s underlying distribution.

    Source: §7 · p. 9

    DINGO’s optimality may not hold across multiple blocks in semi-autoregressive setups, and its constraints are limited to regular languages rather than context-free or context-sensitive languages.

    Source: §6 · p. 9

Abstract

from arXiv · show

Diffusion LLMs have emerged as a promising alternative to conventional autoregressive LLMs, offering significant potential for improved runtime efficiency. However, existing diffusion models lack the ability to provably enforce user-specified formal constraints, such as regular expressions, which makes them unreliable for tasks that require structured outputs, such as fixed-schema JSON generation. Unlike autoregressive models that generate tokens sequentially, diffusion LLMs predict a block of tokens in parallel. This parallelism makes traditional constrained decoding algorithms, which are designed for sequential token prediction, ineffective at preserving the true output distribution. To address this limitation, we propose DINGO, a dynamic programming-based constrained decoding strategy that is both efficient and provably distribution-preserving. DINGO enables sampling of output strings with the highest probability under the model's predicted distribution, while strictly satisfying any user-specified regular expression. On standard symbolic math and JSON generation benchmarks, DINGO achieves up to a 68 percentage point improvement over unconstrained inference

1 Introduction

Diffusion LLMs generate token blocks in parallel, so autoregressive constrained-decoding methods do not directly apply; DINGO addresses the need for valid outputs without distorting their probability distribution.

1 Introduction

  • Constrained decoding matters when outputs feed downstream tools, since syntactic or semantic errors can produce unparsable inputs and cause downstream tasks to fail.The introduction cites logical reasoning with symbolic solvers such as Z3 as one example.Source: §1 · p. 1
  • Diffusion LLMs predict token blocks in parallel, whereas existing constrained decoding relies on sequential token generation to prune invalid choices.The method must also preserve the true output distribution while enforcing constraints.Source: §1 · pp. 1–2
  • The paper frames two challenges: restrict choices so final outputs are structurally correct, and maximize the probability of the generated token block.Source: §1 · p. 2
  • DINGO supports user-specified regular expressions and uses dynamic programming to select the maximum-probability valid string over an output block.The method guarantees that the output remains a valid prefix of some string in the target regular language.Source: §1 · p. 2
  • Experiments report up to a 68% improvement over unconstrained decoding on GSM-symbolic reasoning and JSON generation benchmarks.Source: §1 · p. 2

2 Background

Diffusion LLMs produce output blocks by iteratively unmasking tokens, while regular expressions can be represented by deterministic finite automata for deciding language membership.

2 Background

  • The paper’s notation uses lowercase letters for constants, bold lowercase letters for strings, uppercase letters for functions, and |x| for string length.Source: §2 · p. 2
  • In a single-block setup, the model generates one response block; in a semi-autoregressive setup, it generates multiple blocks sequentially from left to right.Source: §2 · pp. 2–3
  • A diffusion LLM starts with a masked output suffix and iteratively unmasks one or more positions over a fixed number of diffusion steps.The number of steps T is typically predetermined and less than the block length d.Source: §2 · p. 2
  • A DFA for regular expression R processes strings deterministically to decide membership in L(R), using states, a start state, accepting states, and transition rules.Source: §2 · p. 3
  • A live DFA state is one from which some input string can reach an accepting state.Source: §2 · p. 3

3 Optimal Constrained Decoding

The section formalizes constrained decoding for a user-defined regular expression: valid outputs must be extendable to valid completions, and optimal outputs maximize probability among valid strings.

3 Optimal Constrained Decoding

  • A constrained output is correct when it remains a valid prefix of some string in the target regular language.The prefix formulation also accommodates semi-autoregressive generation, where later blocks are appended to the current output.Source: §3 · p. 3
  • For partially masked outputs, validity requires that some replacement of all mask tokens yields a valid prefix of a string in the language.The substitution set contains all fully unmasked strings obtained by replacing each mask token.Source: §3 · p. 3
  • An optimal decoder selects the valid string with the highest probability under the output distribution across response positions.The passage represents each position with a token-probability vector and assumes distributions are independent across positions.Source: §3 · p. 4

4 DINGO Algorithm

DINGO converts a regular-expression automaton for token generation into transitions that handle masks, then uses dynamic programming to find a maximum-probability valid block and establish correctness and optimality.

4 DINGO Algorithm

  • DINGO uses dynamic programming to retrieve an optimal valid output because enumerating the block’s |V|^d candidate strings is impractical.Source: §4 · p. 4

4.1 Precomputation

  • Precomputation builds a token-level DFA by processing each token’s constituent characters, then adds mask transitions to all states reachable by one token transition.The resulting transition function also uses the set of live states.Source: §4.1 · p. 4

4.2 DINGO Dynamic Programming

  • The DP tracks the maximum-probability path to each reachable DFA state, selects the highest-probability path ending in a live state, and reconstructs its token sequence.For semi-autoregressive generation, initialization can start from the live state reached by previously generated blocks.Source: §4.2 · p. 5

4.3 Correctness of DINGO

  • DINGO guarantees a valid prefix and, when one exists for the block length, selects a valid string with probability at least that of any other valid candidate.Source: §4.3 · p. 5

4.4 DINGO algorithm

  • The algorithm’s total complexity is O(d · |Q| · (|Q| + |V|)), combining transition-cost computation and DP updates.The two dominant components are O(d · (|Q|² + |Q| · |V|)) and O(d · |Q|²).Source: §4.4 · p. 6

5 Experiments

DINGO is evaluated on symbolic math and JSON generation, with comparisons against unconstrained generation and a greedy constrained baseline. It improves functional correctness on GSM-Symbolic, achieves perfect syntactic and schema validity on the reported evaluations, and performs well across tested diffusion-block counts.

Experimental Setup

  • The evaluation covers GSM-Symbolic math reasoning and schema-based text-to-JSON generation using LLaDA-8B and Dream-7B base and instruction-tuned models.For math, correctness is checked with Z3; JSON outputs are checked for parsing and schema validity.Source: §5 · pp. 6–8

Baselines

  • The baselines are unconstrained generation, positionwise Greedy Constrained decoding, and a per-problem best-of-greedy-and-unconstrained comparison.Greedy Constrained masks tokens according to the DFA; the best-of baseline selects whichever of the other two methods performs better for each problem.Source: §5 · p. 7

Math Reasoning

  • For LLaDA-8B-I on GSM-Symbolic, DINGO improves functional correctness by 13 percentage points over unconstrained generation and 5 points over Greedy Constrained, with 100% syntactic accuracy across evaluated models.DINGO also introduces marginal overhead compared with unconstrained generation.Source: Table 1 · p. 7, §5 · p. 7

JSON Generation

  • DINGO achieves 100% schema validation and syntactic accuracy on JSON generation, while baseline methods often fail to produce valid JSON.The reported accuracy requires both syntactic validity and validity against the specified schema.Source: Table 2 · p. 8, §5 · p. 8

Ablation Study on The Number of Diffusion Blocks

  • Across 1, 2, and 8 diffusion blocks on GSM-Symbolic, DINGO performs well and outperforms baselines in functional and syntactic correctness.The ablation uses response length 128 and 64 total diffusion steps.Source: §5 · p. 8, Figure 1 · p. 8

6 Related Works

Prior work develops diffusion language models and constrained decoding for autoregressive models, where greedy constraint handling can distort sampling distributions. DINGO applies optimal constrained sampling to diffusion blocks and distinguishes itself through provable constraint guarantees without a surrogate model.

Diffusion Language Models

  • Diffusion LLM research includes semi-autoregressive and block-based approaches that pursue flexible generation, controllability, or inference efficiency.Examples discussed include SSD-LM, LLaDA, BD3-LMs, and Dream-7B.Source: §6 · p. 9

Constrained Decoding with Autoregressive LLMs

  • Autoregressive constrained decoding addresses regular and context-free syntax as well as semantic constraints, but greedy constrained approximations can distort the sampling distribution [Lew et al., 2023, Park et al., 2024a].DINGO instead performs optimal constrained sampling over diffusion-model token blocks, partially mitigating distribution-distortion issues.Source: §6 · p. 9

Constrained Decoding with Diffusion LLMs

  • Concurrent work by Cardei et al. [2025] uses a surrogate-model loss for constrained sampling from diffusion LLMs, but does not guarantee convergence to the constraint and requires a differentiable surrogate.DINGO focuses on provable constraint satisfaction without an additional surrogate model.Source: §6 · p. 9

Limitations

  • DINGO is optimal for per-block generation in pure diffusion settings, but that optimality may not hold in semi-autoregressive multi-block setups.Its current constraints cover regular languages, not context-free or context-sensitive languages or toxicity mitigation.Source: §6 · p. 9

7 Conclusion

The paper concludes that DINGO combines strict regular-language adherence with preservation of the diffusion model’s underlying distribution. Experiments on symbolic math and JSON generation show significant improvements over unconstrained inference for structured-output generation.

Conclusion

  • DINGO uses dynamic programming to enable diffusion LLM outputs to adhere strictly to regular-language constraints while preserving the model’s underlying distribution.The authors report significant improvements over unconstrained inference on symbolic math and JSON generation tasks.Source: §7 · p. 9

NeurIPS Paper Checklist

The NeurIPS Paper Checklist documents authors’ answers and justifications across research practices, clarifies what “NA” means, and describes the diffusion step’s transformer, remasking, and mask-count schedule.

NeurIPS Paper Checklist

  • The authors affirm that the abstract and introduction reflect the paper’s contributions and scope, and report discussing limitations in Related Works.They say the claims are validated by the approach and experiments.Source: NeurIPS Paper Checklist · pp. 13–14
  • The authors answer Yes that the theoretical results have full assumptions and proofs, referring to Section 4.Source: NeurIPS Paper Checklist · p. 14
  • The authors answer Yes that Section 5 provides experimental details needed for reproducibility and understanding results, and state that code is provided.Source: NeurIPS Paper Checklist · pp. 14–16
  • The authors answer No on reporting error bars or other significance information, explaining that greedy decoding makes all experiments deterministic across runs.Source: NeurIPS Paper Checklist · p. 16
  • The authors answer Yes that Section 5 provides compute-resource details and that the research conforms to the NeurIPS Code of Ethics.Source: NeurIPS Paper Checklist · pp. 16–17
  • The authors answer Yes that societal impacts are discussed in the Introduction and evaluated through experiments, and NA for safeguards because they report no such risks.Source: NeurIPS Paper Checklist · p. 17
  • The authors answer Yes that existing assets are credited and new assets documented, citing asset use and the paper’s algorithm and limitations description.Source: NeurIPS Paper Checklist · p. 18
  • The authors answer NA for crowdsourcing and human-subject research, and No for LLM use as an important, original, or non-standard research component.They state that no human subjects were involved and that LLMs were not used for such a component.Source: NeurIPS Paper Checklist · pp. 18–19
  • “NA” means the core method development does not use LLMs as important, original, or non-standard components; the checklist also points to the NeurIPS LLM policy.Source: NeurIPS Paper Checklist · p. 19

A Transformer and Remasking Step

  • At each diffusion step, the transformer computes token distributions, then remasking sets selected positions to the mask token with probability 1 and all other tokens to probability 0.The remasking strategies described include random positions and positions chosen by lowest top-token probability (Nie et al. [2025]), or highest entropy (Ye et al. [2025]).Source: A Transformer and Remasking Step · p. 20
  • The number of masked positions starts at the block length d and decreases linearly to zero over T diffusion steps.Source: A Transformer and Remasking Step · p. 20

B Proofs

The proofs establish DINGO’s valid-prefix guarantee and maximum-probability choice among valid block strings. For regular expressions with a small constant number of DFA states, parallelized DP has linear depth in vocabulary size and block length, while semi-autoregressive generation carries DFA state between blocks.

B Proofs

  • If at least one fully unmasked valid prefix has nonzero probability, DINGO retrieves a valid string that is a prefix of a target-language string.The proof fills masked positions with tokens that preserve DFA reachability, yielding a string in the target language’s prefix set.Source: B Proofs · pp. 20–21
  • DINGO’s decoded string has probability at least as high as any other valid string under the model’s block output distribution.The proof identifies the terminal DP score with the decoded string’s probability and bounds other valid strings by that score.Source: B Proofs · pp. 21–22

C Time complexity analysis of parallelized DINGO DP

  • O(|V| + d) computational depth for regular expressions with a small constant number of DFA states, linear in vocabulary size and block length.The stated general depth is O(max(|Q|^2, |Q| × |V|) + |Q|^2 × d).Source: C Time complexity analysis of… · p. 22

D Semi-Autoregressive

  • In semi-autoregressive decoding, each block carries forward the current DFA state, and generation returns None if the output so far is not in a live state.The carried state initializes the next block’s dynamic programming step.Source: D Semi-Autoregressive · p. 23

E Token Transitions Statistics

The appendix specifies symbolic-expression and JSON tasks, their regular-expression formats, and examples contrasting unconstrained, greedy constrained, and DINGO outputs. For transition precomputation, the reported maxima are 455 DFA states and 17.7 seconds for Dream and 21.3 seconds for LLaDA.

E Token Transitions Statistics

  • The maximum DFA size is 455 states, and maximum precomputation times are 17.7 seconds for Dream and 21.3 seconds for LLaDA.For JSON generation, the reported precomputation time and state count are means across schemas.Source: E Token Transitions Statistics · p. 24

F.1 GSM-Symbolic Prompt

  • The GSM-Symbolic prompt asks for a step-by-step response with intermediate expressions in << >> and a final symbolic expression using specified operations.The prompt examples include arithmetic word problems with variables.Source: §F.1 · p. 24

F.2 GSM-Symbolic Regex

  • The GSM-Symbolic regex listing supplies a regular expression for the task’s permitted symbolic-expression output.The listing follows the task prompt’s expression format.Source: §F.2 · p. 25

F.3 GSM-Symbolic Case Studies

  • In the GSM-Symbolic examples, unconstrained output is syntactically incorrect, greedy constrained output is valid but incorrect, and DINGO produces the correct answer.The figures identify the examples as GSM-Symbolic cases.Source: Figure 2 · p. 25, Figure 3 · p. 26

G.1 JSON-Mode Example Prompt

  • The JSON example prompt supplies a schema and asks for a JSON object that follows it, using specified campaign details.The example schema includes campaign, product, date, and discount fields.Source: §G.1 · p. 26

G.2 JSON-Mode Example Regex

  • The JSON example regex encodes the required campaign and product identifiers, dates, and optional discount-details field.The regex is presented as the format for the example JSON schema.Source: §G.2 · p. 26

G.3 JSON-Mode Case Studies

  • In the JSON example, unconstrained output is syntactically incorrect, greedy constrained output is valid but incomplete, and DINGO generates a syntactically correct answer.The figures identify the example as JSON generation.Source: Figure 4 · p. 27, Figure 5 · p. 27

H Ablation Study on Number of Blocks for Diffusion LLM Generation (GSM-Symbolic)

The GSM-Symbolic block ablation uses a response length of 128, 64 diffusion steps, and 1, 2, or 8 blocks; the supplied table caption identifies its ablation scope.

H Ablation Study on Number of Blocks for Diffusion LLM Generation (GSM-Symbolic)

I Ablation Study on Number of Blocks for Diffusion LLM Generation (JSON-Mode)

The JSON-Mode block ablation uses a response length of 128, 64 diffusion steps, and 1, 2, or 8 blocks; the supplied table caption identifies its ablation scope.

I Ablation Study on Number of Blocks for Diffusion LLM Generation (JSON-Mode)

J Ablation Study on Number of Steps for Diffusion LLM Generation (GSM-Symbolic)

The supplied ablations vary diffusion-step counts for GSM-Symbolic and JSON-Mode with Dream-I-7B; both table captions identify the task and model.

J Ablation Study on Number of Steps for Diffusion LLM Generation (GSM-Symbolic)

K Ablation Study on Number of Steps for Diffusion LLM Generation (JSON-Mode)

Related papers

Loading 2505.23061v1…