PaperlayerView exact revision on arXiv

Source-linked AI summary

Accelerating Diffusion Language Models via Structured Suffix Modeling

Zifeng Cheng, Keda Li, Zhiwei Jiang, Cong Wang, Fei Shen, Qing Gu

arXiv:2608.23167v1cs.CL

TL;DR

DLM inference incurs substantial suffix computation, and existing locality-based pruning overlooks regional suffix roles and previous-step denoising information. SSM models local, middle, and tail suffix regions differently while carrying denoising information across steps, reducing latency and achieving up to a 72.81× speedup in long-sequence inference when combined with other acceleration techniques.

  • Problem

    DLM inference incurs substantial suffix computation, while existing locality-based pruning overlooks differing suffix-region roles and discards previous-step denoising information.

  • Method

    SSM is a training-free framework that assigns different token budgets to local, middle, and tail suffix regions and incorporates previous-step decoding into current suffix representations.

  • Results

    Across mathematical reasoning and code generation benchmarks on three DLMs, SSM reduces inference latency while maintaining competitive accuracy, achieving up to a 72.81× speedup in long-sequence inference with other acceleration techniques.

  • Takeaways & Limitations

    SSM offers inference-time acceleration that complements parallel decoding and KV cache while maintaining competitive or improved task performance.

  • Takeaways & Limitations

    SSM may degrade performance on some datasets, limiting its suitability for scenarios where accuracy is highly critical.

Abstract

from arXiv · show

Diffusion Language Models (DLMs) exhibit strong parallel decoding capabilities by denoising multiple tokens in a single generation step. However, this parallelism comes with substantial computational overhead, as each step requires interactions with all suffix tokens. Existing methods typically reduce this cost by retaining only a local suffix window as a substitute for the full suffix. Despite their effectiveness, these methods overlook the structural heterogeneity across suffix regions and re-initialize suffix tokens with identical representations at each timestep. To this end, we propose a structured suffix modeling method for efficient DLM inference. Specifically, we divide the suffix into three regions, i.e., the local, middle, and tail regions, and retain different numbers of suffix tokens in each region according to their structural roles. Moreover, we incorporate the decoding results from the previous step into the suffix token representations at the current step, allowing them to carry evolving denoising information across generation steps. Notably, our method is training-free and orthogonal to several existing acceleration techniques, such as parallel decoding strategies and KV cache. Empirical results across multiple benchmarks on three DLMs demonstrate that our method can further accelerate DLM inference and improve performance in most cases. In particular, in long-sequence inference, our method achieves up to a \(72.81\times\) speedup when combined with other acceleration techniques. Our code is available at https://github.com/zifengcheng/SSM.

1 Introduction

Diffusion language models enable parallel decoding but incur suffix-related computational overhead and overlook regional suffix roles and evolving denoising information. SSM addresses these limitations with structured suffix modeling and soft suffix embeddings, providing training-free, composable inference acceleration with competitive task performance.

  • DLMs use bidirectional context to predict multiple tokens per generation step, reducing latency relative to autoregressive decoding but requiring suffix-token interactions during denoising.The supplied passage introduces parallel decoding over entire sequences or blocks and identifies suffix interactions as a source of computational overhead.
  • Existing suffix-dropout methods treat suffix pruning primarily as a locality problem, overlooking distinct contextual, structural, and boundary roles across suffix blocks.The introduction also notes that suffix tokens are re-initialized at each timestep without carrying information from earlier decoding steps.
  • SSM divides the suffix into local, middle, and tail regions, assigning structured modeling based on their different roles during block-wise decoding.The regional division is motivated by attention scores that rapidly decrease, stabilize, and then increase across suffix positions.
  • SSM is training-free, applies directly to existing DLMs, and is orthogonal to parallel decoding and KV cache, reducing inference latency while maintaining competitive accuracy.Experiments cover mathematical reasoning and code generation benchmarks using LLaDA-Instruct, LLaDA-1.5, and Dream-Base.
  • SSM’s soft suffix embedding incorporates previous-step decoding results into current suffix representations, allowing denoising information to evolve across generation steps.This mechanism is designed to replace timestep-wise re-initialization with representations that carry information forward.

2 Related work

Diffusion language models accelerate generation through parallel denoising, motivating decoding-strategy and suffix-reduction methods. Suffix dropout removes tokens without accessing internal states, while this work instead targets small suffix-token settings without iterative refinement.

  • Accelerating Diffusion Language Models: dLLMs use bidirectional attention to denoise multiple tokens per timestep, motivating methods that accelerate inference and exploit parallel decoding.
  • Accelerating Diffusion Language Models: Decoding-strategy methods define new criteria, including Fast-dLLM’s confidence-threshold decoding and latent refinement decoding.
  • Suffix Dropout: Suffix dropout accelerates DLM inference by excluding suffix tokens in advance, unlike sparsification methods that require internal model states such as attention weights.
  • Suffix Dropout: DPad models suffix tokens as a non-semantic reservoir, using a local window with distance-aware dropout at each block’s decoding start.
  • Suffix Dropout: Streaming-dLLM preserves neighboring suffix tokens and the final token to maintain information during inference.

3 Method

SSM is a training-free framework that treats suffix tokens as a structured information field, allocating region-specific token budgets to reduce redundant computation while preserving contextual and boundary information. It further propagates previous-step denoising information through soft suffix embedding mixing and supports sparse inputs, positional preservation, and early termination for efficient DLM inference.

  • Structured Suffix Modeling: SSM divides the suffix into local, middle, and tail regions because their attention patterns and functional roles differ.The local region provides fine-grained context, the middle region lightweight structural anchors, and the tail region terminal boundary information.
  • Suffix Region Partition: The first w suffix blocks form the local region, the last suffix block forms the tail region, and intervening blocks form the middle region.If fewer than w suffix blocks exist, all suffix blocks are assigned to the local region.
  • Suffix Token Retention: SSM retains all local tokens, only each middle block’s start token, and each tail block’s start and final tokens.The middle start token preserves skeletal structure, while the tail final token provides boundary information for planning future generation length.
  • Soft Suffix Embedding Mixing: Soft suffix embedding mixing combines previous-step decoded-token information with [MASK] representations so suffix tokens carry evolving denoising information across steps.The previous prediction is represented by a weighted combination of top-k token embeddings before being mixed into the current suffix embedding.
  • Sparse Input Construction: SSM constructs sparse inputs from the prompt, current block, and retained suffix tokens while preserving their original positional encodings; it can also stop generation after detecting <eos>.Early termination is especially useful when maximum generation length substantially exceeds the actual output length.

4 Experiments … 4.3 Main Results

Across four reasoning and code-generation benchmarks on three diffusion language models, SSM generally improves efficiency while maintaining competitive accuracy. Its gains come from structured suffix reduction and more concise generations, though TPS can diverge from end-to-end latency.

  • 4.1 Experimental Settings: Experiments cover GSM8K, MATH, HumanEval, and MBPP, using flexible-extract, strict-match, pass@1, average latency, TPS, and ¯ℓ/ℓmax.Mathematical tasks use flexible-extract and strict-match; code tasks use pass@1.
  • 4.1 Experimental Settings: SSM is evaluated on LLaDA-8B-Instruct, LLaDA-1.5, and Dream-7B-Base using A800 GPUs, block size 32, and confidence threshold 0.9.The search ranges include top-k values [3, 5, 7] and mixing ratios α of [0.2, 0.3, 0.4].
  • 4.2 Baselines: The baselines include Vanilla top-1 decoding and threshold-based parallel decoding, with SSM combined with parallel decoding for comparisons against DPad+Par. and Streaming+Par.DPad applies distance-aware dropout in a local window, while Streaming preserves neighboring suffix tokens and the final token.
  • 4.3 Main Results: SSM outperforms baseline methods in most cases, achieves the lowest latency, and delivers a 2.3× speedup over Streaming+Par. on MBPP with LLaDA-Instruct.Compared with Par. alone, SSM reduces latency by 43.7%, 38.9%, 71.1%, and 81.4% on LLaDA-Instruct.
  • 4.3 Main Results: SSM improves efficiency by assigning different token budgets to suffix regions and often producing shorter generations with smaller ¯ℓ/ℓmax values.The shorter outputs are attributed to reduced suffix context and less redundant information, rather than early termination.
  • 4.3 Main Results: Latency and TPS can diverge because substantially shorter generations may reduce GPU utilization and mechanically lower TPS despite better end-to-end latency.This limitation motivates more appropriate throughput metrics for future evaluation (Qian et al., 2026; Chen et al., 2026).
  • 4.3 Main Results: SSM maintains competitive accuracy and often improves over strong baselines, with larger gains on code tasks and mathematical strict-match.Strict-match also reflects whether the model follows the demonstrated reasoning format, providing an indication of in-context learning ability.

4.4 Ablation Study

The ablation study finds that every proposed component contributes to performance improvement. Removing the local region causes severe degradation and unusually short generations, while previous-step suffix mixing improves efficiency and performance under dataset-dependent effects.

  • 4.4 Ablation Study: All proposed components contribute to the observed performance improvement.Table 2 evaluates the contribution of each component.
  • 4.4 Ablation Study: Removing the local region severely degrades performance and shortens average generation length, showing that neighboring suffix tokens provide essential fine-grained context.The local region should therefore be densely retained to prevent premature or incomplete generation.
  • 4.4 Ablation Study: Removing middle-region start tokens hurts accuracy on both datasets and increases latency, supporting their retention.
  • 4.4 Ablation Study: Soft suffix embedding mixing reduces latency and improves GSM8K performance, while substantially shortening generation length and reducing latency on HumanEval.These effects indicate that incorporating previous-step decoding results into current suffix representations preserves evolving denoising information.
  • 4.4 Ablation Study: Early termination removes time spent generating subsequent blocks after EOS detection without changing generation length.

4.5 Effects of Hyper-Parameters · 4.6 Comparison with Other Parallel Decoding Methods

The method’s hyper-parameters expose dataset-dependent trade-offs, with α = 0.4 giving the best accuracy before overly weighting unreliable early predictions becomes harmful. Against margin-based parallel decoding, it achieves the best accuracy and latency and combines orthogonally for further acceleration.

  • 4.5 Effects of Hyper-Parameters: Increasing local-region size raises latency on both datasets because retaining more suffix tokens adds computation.On GSM8K, performance generally improves with a larger local region, indicating that local-region information is particularly beneficial there.
  • 4.5 Effects of Hyper-Parameters: HumanEval is more sensitive to local-region size than GSM8K, according to the reported hyper-parameter analysis.The supplied passage begins describing this sensitivity but does not provide the complete comparison.
  • 4.5 Effects of Hyper-Parameters: For α values from 0.2 to 0.4, latency and accuracy remain relatively stable across both datasets, with best accuracy at α = 0.4.The result holds for the tested α settings before the sharp accuracy decline at larger values.
  • 4.5 Effects of Hyper-Parameters: Increasing α to 0.5 or 0.6 sharply reduces accuracy because unreliable early-timestep predictions receive overly large weights and corrupt suffix representations.The explanation is presented as the likely cause of the observed degradation.
  • 4.6 Comparison with Other Parallel Decoding Methods: The comparison evaluates a margin-based strategy that decodes tokens when the top-two prediction probability gap exceeds a predefined threshold.The strategy is attributed to Kim et al. (2025a).
  • 4.6 Comparison with Other Parallel Decoding Methods: Our method achieves the best accuracy and latency against three margin-based parallel decoding baselines, while also accelerating that strategy by 3.33× on HumanEval.This demonstrates that suffix dropping is orthogonal to parallel decoding and can improve speed and performance when combined with it.

4.7 Effects of Generation Length · 4.8 Performance on Long-sequence Generation

SSM maintains lower latency across generation lengths, with larger gains for longer outputs, and remains effective on long-sequence generation when combined with parallel decoding, prefix caching, and KV cache. On GSM8K with LLaDA-1.5, SSM improves accuracy in most settings and reaches a 72.81× speedup over Vanilla top-1 decoding with parallel decoding and prefix caching.

  • 4.7 Effects of Generation Length: SSM achieves lower latency than DPad at every tested generation length, with the advantage increasing as generation length grows.Figure 4 compares the methods on two datasets and shows that shorter generation lengths generally produce lower latency.
  • 4.7 Effects of Generation Length: On HumanEval with generation length 1024, SSM is nearly 3× faster than DPad.The comparison covers different generation lengths on two datasets, and the method accelerates inference across these lengths.
  • 4.8 Performance on Long-sequence Generation: SSM consistently has lower latency than DPad and better accuracy in most GSM8K settings across acceleration-technique combinations.The evaluation uses LLaDA-1.5 with a 1024-token maximum generation length and compares methods under matched settings.
  • 4.8 Performance on Long-sequence Generation: Long-sequence generation provides larger opportunities for suffix-dropout acceleration because it involves longer suffixes.The long-sequence evaluation uses GSM8K, LLaDA-1.5, and a maximum generation length of 1024 tokens.
  • 4.8 Performance on Long-sequence Generation: SSM remains compatible with KV cache and parallel decoding, enabling further speedup beyond either technique alone.This confirms that SSM is orthogonal to decoding-based and cache-based acceleration methods.

5 Conclusion

SSM is a training-free framework that accelerates DLM inference through structured suffix modeling, allocating different token budgets across suffix regions and incorporating prior decoding information into current suffix representations.

  • SSM accelerates DLM inference with a training-free structured suffix modeling framework.
  • Different token budgets for local, middle, and tail suffix regions preserve essential contextual, structural, and boundary cues while reducing redundant suffix computation.
  • Soft suffix embeddings incorporate previous-step decoding information into current suffix representations.

A Limitations

The method’s efficiency gains are limited when the suffix is a small part of the sequence, and it can reduce performance on some datasets. Its behavior on models larger than 10B parameters remains unevaluated.

  • A Limitations: Efficiency improvements are relatively limited when the suffix constitutes a small proportion of the sequence, such as with a long prefix or short generation length.
  • A Limitations: The method may degrade performance on some datasets, making it less suitable when accuracy is highly critical.
  • A Limitations: The method has not been evaluated on models larger than 10B parameters.

B Results on LLaDA-1.5

On LLaDA-1.5, SSM+Par. validates structured suffix modeling across four benchmarks, delivering stable efficiency gains while maintaining or improving task accuracy. It achieves the strongest reported GSM8K accuracies and substantially improves strict-format performance over Vanilla.

  • The results further validate SSM’s effectiveness on the advanced LLaDA-1.5 diffusion language model across four benchmarks.
  • Efficiency: SSM+Par. delivers superior, stable acceleration across all four benchmarks and consistently achieves lower latency than DPad and Streaming.The reported efficiency is accompanied by shorter average generated sequences, reaching 73/512 on HumanEval and 49/512 on MBPP.
  • Accuracy: 80.14% Strict accuracy on GSM8K, up from 61.87% with Vanilla, while achieving the benchmark’s highest Flexible and Strict scores.SSM+Par. can maintain task performance and sometimes improve it, with especially pronounced gains relative to LLaDA-8B-Instruct results.

C Comparison under Vanilla Top-1 Decoding

Under vanilla Top-1 decoding, SSM achieves the best latency and accuracy in most cases, demonstrating effectiveness across decoding strategies.

  • C Comparison under Vanilla Top-1 Decoding: SSM achieves the best latency and accuracy in most cases under vanilla Top-1 decoding.The comparison uses LLaDA-Instruct across three benchmarks.
  • C Comparison under Vanilla Top-1 Decoding: The results show that SSM remains effective when evaluated with Top-1 decoding.
  • C Comparison under Vanilla Top-1 Decoding: SSM is compatible with different decoding strategies, including vanilla Top-1 decoding.

NeurIPS Paper Checklist

The checklist states that the paper clearly presents its claims, documents limitations, relies on empirical rather than theoretical validation, and provides experimental details and code for reproducibility. It also explains that error bars are omitted because the training-free inference setting has no random initialization or training-run variance.

  • Claims and scope: The paper states that its abstract and introduction clearly describe the claims and contributions.
  • Limitations: A limitations section is provided in the appendix.
  • Theory assumptions and proofs: The paper supports its claims with extensive empirical results rather than theoretical proofs.
  • Experimental reproducibility and setting: The authors state that hyperparameters and experimental details are provided to reproduce and understand the main results.These details are reported in the Experimental Settings.
  • Open access to data and code: The authors provide data and code on GitHub.
  • Experiment statistical significance: No error bars or statistical significance tests are reported because the training-free, fixed inference-only evaluation has no variance from random initialization or training runs.
Loading 2608.23167v1…