GitHub

S2D2: Fast Decoding for Block-Diffusion LLMs via Training-Free Self-Speculation

Code for the paper.

Overview

Block-diffusion language models offer a promising path toward faster-than-autoregressive generation by combining block-wise autoregressive decoding with within-block parallel denoising. However, in the few-step regime needed for practical acceleration, standard confidence-thresholded decoding is often brittle: aggressive thresholds hurt quality, while conservative thresholds require unnecessary denoising steps. Existing approaches that address this issue either require additional training or incur extra test-time compute.

We present S2D2, a training-free self-speculative decoding framework for block-diffusion language models. Our key observation is that a block-diffusion model becomes autoregressive when the block size is reduced to one, allowing the same pretrained model to act as both drafter and verifier. S2D2 inserts a speculative verification step into standard block-diffusion decoding and uses lightweight routing policies to decide when verification is worth its cost. This yields a hybrid decoding trajectory in which diffusion proposes tokens in parallel, while the autoregressive mode acts as a local sequence-level critic.

Across three mainstream block-diffusion families, S2D2 consistently improves the accuracy-speed tradeoff over strong confidence-thresholding baselines. On SDAR, we observe up to 4.7x speedup over autoregressive decoding, and up to 1.57x over a tuned dynamic decoding baseline while improving accuracy by up to 4.5 points. On LLaDA2.1-Mini, S2D2 remains complementary to built-in self-correction, including a conservative setting where it is 4.4x faster than the static baseline with slightly higher accuracy.

Project Structure

This codebase covers four block-diffusion model families. For each model, we copy the core functionality files into separate subfolders with our modifications. To run experiments, cd into the corresponding subfolder.

S2D2/
├── SDAR/                  # SDAR-8B-Chat
├── Fast-dLLM-v2/          # Fast-dLLM v2
├── LLaDA2/                # LLaDA2.1-Mini
└── D2F/                   # Discrete Diffusion Forcing

Installation

Please follow each model's official instructions to install required packages. We highlight the following version requirements:

  • SDAR requires transformers==4.52.4. Flash Attention must be installed as:
    pip install "flash-attn==2.7.4.post1" --no-build-isolation --no-cache-dir
  • Other models generally work with a wider range of dependency versions.

Routing Policies

The following arguments control S2D2's routing policies and are shared across all example and evaluation scripts:

Argument Description
--do_verify_policy Routing policy: mask_span_length, score_threshold, score_hysteresis, contextual_bandit_ucb
--do_verify_score_threshold Score threshold

Read the original on github.com ↗