Source-linked AI summary
Dystruct: Dynamically Structured Diffusion Language Model Decoding via Bayesian Inference
Bian Sun, Kevin Zhai, Mubarak Shah, Zhenyi Wang
TL;DR
DLMs offer parallel decoding but commonly use fixed lengths, while existing flexible-length methods may require retraining or rely on local criteria that overlook sequence structure. DyStruct introduces training-free Bayesian structured decoding that jointly infers expansion size, block partitions, and decoding order, and experiments report consistent improvements over fixed-length and flexible-length baselines.
Problem
Most DLMs use fixed generation lengths, while existing flexible-length methods require retraining or rely on local confidence signals that overlook evolving sequence structure.
Method
DyStruct performs training-free Bayesian joint inference over window expansion size, contiguous block partitions, and block decoding schedules, using a CRP prior for adaptive structure.
Results
DyStruct consistently outperforms fixed-length and existing flexible-length DLM decoding methods across diverse datasets and benchmarks.
Takeaways & Limitations
Bayesian structured decoding provides a principled inference-time mechanism for flexible-length, coherent generation while keeping the model frozen.
Takeaways & Limitations
The method operates only at inference time; integrating structural inference into training may further improve performance.
Abstract
from arXiv · showhide
Diffusion language models (DLMs) have recently emerged as a promising alternative to autoregressive models, primarily due to their ability to enable parallel decoding. Despite this advantage, most existing DLMs rely on a fixed generation length specified prior to decoding, which restricts their flexibility in real-world applications. While a few recent works attempt to support flexible-length generation, they typically suffer from notable limitations: some require costly retraining to accommodate variable-length outputs, while others depend solely on local confidence signals during decoding. Such local criteria fail to capture the evolving structure of the sequence, often resulting in suboptimal generation quality. In this paper, we propose a training-free, Bayesian structured decoding framework that formulates flexible-length generation as a dynamic structural inference problem. Our approach formulates flexible-length generation as a dynamic structural inference problem, jointly computing the expansion length, the block boundaries, and the decoding schedule. At each window expansion step, the method integrates local uncertainty with structural signals via a unified mechanism that supports dynamic structured generation, including both flexible block expansion and block organization, while maintaining coherence. Extensive experiments across multiple benchmarks demonstrate that our approach significantly improves generation quality and flexibility over existing fixed-length and flexible-length baselines. These results highlight the advantage of Bayesian structured decoding for diffusion language model, providing a principled and efficient solution for structured text generation.
1 Introduction
Diffusion language models enable parallel decoding but usually require a fixed generation length, limiting adaptation to task complexity and evolving context. DyStruct addresses this with training-free Bayesian structured decoding that jointly infers expansion length, block structure, and decoding order.
- Motivation: Fixed-length DLM decoding can truncate complex responses or introduce redundancy in simple ones.The optimal output length depends on task complexity, motivating dynamic adjustment during generation.
- Limitations of prior work: Existing flexible-length approaches either require costly retraining or rely on local confidence heuristics that overlook post-expansion content organization.The latter can produce fragmented structure when new tokens are generated.
- Method: DyStruct formulates flexible-length generation as Bayesian joint inference over window expansion size, contiguous block partitions, and block decoding schedule.The framework uses a structured prior over latent block partitions to guide content organization during decoding.
- Method: A Chinese Restaurant Process prior adaptively determines block quantity and boundaries while providing predictive guidance for continuing or starting blocks.This avoids preset block counts and predefined partition boundaries.
- Method: The method performs structured decoding at inference time without modifying model parameters or requiring additional training.Its schedule is dynamically determined using block instability and context-aware prioritization.
4 Method
DyStruct performs flexible-length diffusion decoding by jointly inferring window size, contiguous block partitions, and block decoding order from diagnostic observations and structured priors. Its procedure combines adaptive expansion, CRP-based partitioning, context-aware scheduling, and iterative refinement.
- DyStruct models flexible-length decoding as Bayesian inference over window length, block partition, and decoding schedule.
- The next window length adapts to preceding-window instability, with unstable windows restricting expansion and stable windows permitting larger expansions.
- A temporary diagnostic pass extracts token-level signals, including entropy, prediction shifts, hidden-state variation, and confidence, to estimate local instability.
- The CRP prior partitions each window at inferred gaps into contiguous blocks, favoring coherent block growth while allowing dynamically supported new boundaries.
- A schedule prioritizes anchored blocks with low instability so decoded tokens provide stable context for subsequently decoding more unstable regions.
- Within each scheduled block, high-confidence tokens are committed while remaining masked positions are iteratively refined, followed by localized edge-welding at shared boundaries.
5 Experiments
DyStruct is evaluated across mathematical, coding, and logical reasoning benchmarks under matched computational budgets, with additional ablations examining efficiency and structural components. Results show improved benchmark accuracy, adaptive compute allocation, and sensitivity to scheduling, partitioning, and initial-window choices.
- 5.1 Main Results: DyStruct improves accuracy across all five benchmarks, including BBH exact match from 44.9 to 49.3 on LLaDA-8B.Paired McNemar tests found significant prompt-level improvements for BBH and both mathematics datasets.
- 5.1 Main Results: DyStruct raises MBPP accuracy from 39.8 to 41.4 on LLaDA-8B while transferring consistently to Dream-7B.The evaluation uses greedy pass@1 accuracy for code generation and strict computational controls.
- 5.1 Main Results: DyStruct terminates refinement early on low-instability GSM8K regions, producing lower seconds-per-iteration inference time across both backbones.On BBH, it instead allocates computation toward high-instability logical transitions, yielding a 4.4-point accuracy improvement.
- 5.2 Structural Ablations and Sensitivity: Removing context-aware Gibbs scheduling reduces MATH accuracy from 31.4 to 30.3, while removing edge-welding lowers HumanEval pass rates by 1.9 points.The ablation attributes these degradations to fixed left-to-right ordering and incompatible syntax across independently decoded block boundaries.
- 5.3 Qualitative Analysis: An expansion length of 48 tokens maximizes BBH accuracy, while average expansion terminates at 219 HumanEval tokens and 246 BBH tokens below the 256-token limit.Performance remains stable across concentration-prior variations, and fixed structural hyperparameters work across evaluated datasets without task-specific tuning.
- 5.3 Qualitative Analysis: The qualitative analyses show edge-welding repairs entropy-spiking boundaries, partitioning isolates high-instability deductions, and scheduling anchors stable terminal blocks before intermediate inference.Figures 3–5 connect these operations to syntax repair, logical-transition refinement, and bidirectional conditioning.
6 Conclusion
DyStruct provides a Bayesian framework for flexible-length diffusion language model decoding and reports consistent gains over fixed-length and existing flexible-length methods. Its inference-only design leaves integrating structural inference into training as future work.
- DyStruct consistently outperforms fixed-length and existing flexible-length DLM decoding methods across multiple benchmarks.
- The method operates purely at inference time without modifying model parameters.This enables broad applicability.
- Integrating structural inference into training may further enhance performance, but the paper leaves this direction for future work.
A Notation
The notation appendix provides a mathematical reference for DyStruct's notation, presented in Table 5.
- Table 5 provides a concise mathematical reference for the notation used in DyStruct.
B Algorithm
Algorithm 2 executes DyStruct through adaptive expansion, diagnostic partitioning, instability-aware scheduling, block-wise resolution, edge welding, and state updates until generation ends.
- Initialization and loop: Algorithm 2 initializes the prompt, length limit, hyperparameters, sequence, step, and previous-window instability before iterating until EOS or the maximum length.
- Adaptive Window Expansion: At each iteration, the decoder samples an expansion length from a Poisson distribution scaled by one minus previous instability, then appends masked tokens.
- Diagnostic Pass & Partitioning: A temporary diagnostic pass extracts feature signals and computes position-instability and gap-split probabilities for partitioning.
- Instability-Aware Scheduling: Blocks are ordered by sorting negative mean instability plus context adjacency, producing an instability-aware decoding schedule.
- Block-wise Resolution & Edge-Welding: Each scheduled block is refined and committed, while low-confidence positions near shared boundaries are remasked and refined through edge welding.
- State Update: The algorithm updates the sequence and iteration state after calculating finalized mean instability over the decoded window, then outputs the final sequence.
- Algorithmic Description: DyStruct determines new-window length from prior-segment stability and uses CRP-based Bayesian partitioning to group tokens into contiguous blocks.
C Method Details
For each gap between adjacent tokens, DyStruct constructs a feature vector from diagnostic distributions and Jensen-Shannon divergence to assess block continuity.
- For each internal gap, DyStruct forms a feature vector using diagnostic information around the adjacent token positions.
- Jensen-Shannon divergence helps quantify whether adjacent tokens form one contiguous block or require structural separation.
D Calibration of Instability Coefficients
DyStruct calibrates token instability from seven diagnostic features and binary ground-truth targets, then projects and normalizes the learned score during inference.
- The calibration dataset contains token-level observations extracted from validation trajectories with accessible ground-truth sequences.
- Each feature vector ϕ_i contains predictive entropy, remasking frequency, logit oscillation, Jensen-Shannon divergence, hidden-state jump, confidence, and probability margin.
- The binary target d_i marks mismatches or remasking events as unstable and correctly committed tokens as stable.
- The feature components are computed from logits, hidden states, probabilities, predictions, masking indicators, and token-acceptance indicators.
- The coefficient vector w* is estimated by minimizing an L2-regularized binary cross-entropy loss over the calibration data.
- During inference, a linear projection preserves calibrated instability rankings, and a window-centered logistic function converts it into the final positional score h_i.
E Implementation Details and Hyperparameters
DyStruct uses fixed implementation settings for expansion, Bayesian partitioning, scheduling, block refinement, and edge welding across all benchmarks and model scales.
- DyStruct maintains one uniform hyperparameter set across benchmarks, domains, and model scales without task-specific tuning.
- Window expansion is bounded by a minimum burst length of Lmin = 8 and a maximum burst length of Lmax = 48.
- Bayesian partitioning uses a CRP concentration prior of α0 = 1.5, while schedule evaluation assigns context adjacency a priority weight of γ = 2.0.
- Block refinement interpolates between Tmin = 6 and Tmax = 18 based on instability, and welding uses rweld = 4 tokens for 4 refinement steps.
F Ablative Studies
The welding-radius ablation shows that moderate local repair preserves block structure and boundary coherence, whereas overly large radii can weaken locality and degrade performance.
- The welding interval uses a radius-controlled neighborhood around each block boundary to reconcile adjacent predictions without changing the partitioning objective.
- Moderate welding radii provide enough cross-boundary context to correct local inconsistencies while preserving the original block structure.
- On HumanEval, results remain unchanged for rweld = 8 to rweld = 12, while nearby MBPP settings remain close to the default rweld = 10.
- Overly small radii may leave adjacent blocks weakly aligned by excluding sufficient cross-boundary dependency.
- Overly large radii may remask stable positions, weaken Bayesian locality, and slightly degrade HumanEval performance at rweld = 16.
- The main experiments therefore use rweld = 10 as a fixed default for stable boundary reconciliation without task-specific tuning.
G Statistical Significance
Paired statistical tests show that DyStruct improves over DAEDAL across BBH, Math, and Code evaluation groups, with the strongest evidence for BBH and Math.
- The paired McNemar comparison evaluates discordant prompts solved by only one of DyStruct and DAEDAL on the same prompt sets.
- DyStruct consistently improves over DAEDAL across all three evaluation groups in the paired McNemar analysis.
- BBH shows a substantially higher DyStruct accuracy than DAEDAL, with highly significant differences under both corrected McNemar and exact binomial tests.
- Pooled Math benchmarks also show a statistically significant improvement, extending the observed gains beyond BBH reasoning tasks.
- Code accuracy is higher with DyStruct, but the smaller 664-instance evaluation group provides lower statistical power and supports only a positive trend.