PaperlayerView exact revision on arXiv

Source-linked AI summary

Early-Bird Decoding: Accelerating Diffusion LLMs with Learnable Block Sizes and Parallel Sampling

Lixuan Wei, Wei Zhou, Jianwen Wu, Yipeng Shen, Meiling Wang, Haoran You

arXiv:2609.16450v1cs.CLcs.AIcs.LG

TL;DR

dLLMs promise parallel generation, but fixed blocks and confidence thresholds can leave inference inefficient. EB-Decode learns variable-length token groups and early commitment while keeping the base model frozen; across three open-source dLLMs and four benchmarks, it reports up to 18.76× throughput over vanilla decoding with comparable accuracy.

  • Problem

    Fixed block sizes and confidence thresholds do not adapt to heterogeneous token difficulty, limiting dLLM inference efficiency.

    Source: §1 · pp. 1–2

  • Method

    EB-Decode uses learnable routers for variable-length token grouping and position-aware early commitment without modifying the base dLLM weights.

    Source: §1 · p. 2, §5.2 · p. 7, §7 · p. 12

  • Results

    Across three open-source dLLMs and four benchmarks, EB-Decode achieves up to 18.76× throughput improvement over the vanilla decoder with comparable accuracy.

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

  • Takeaways & Limitations

    The results support EB-Decode as a plug-in approach to accelerating dLLM inference while keeping pretrained weights frozen.

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

    Evaluation uses generation lengths of 256 and 512 tokens and batch size 1; substantially longer lengths and realistic serving conditions remain future work.

    Source: I Limitations · p. 22

Abstract

from arXiv · show

Diffusion large language models (dLLMs) offer a promising parallel decoding paradigm as an alternative to autoregressive generation through iterative unmasking. However, dLLMs typically require many steps before token confidence reaches the decoding threshold, resulting in inefficient inference even with block-wise KV caching. To accelerate dLLM inference, we for the first time propose an "early-bird (EB)" decoding framework, motivated by the observation that tokens with similarly low entropy tend to cluster and can be jointly decoded earlier, before reaching the confidence threshold. In particular, our EB-Decode framework integrates two key enablers: (1) a learnable network that adaptively groups tokens with similar uncertainty into variable-length blocks, rather than relying on fixed block sizes; (2) a position-aware sampler that learns to unmask tokens in parallel using fewer decoding steps within predicted variable-length blocks. Both components are developed without modifying pretrained dLLM weights and can therefore be directly deployed as plug-ins during serving, with negligible training and inference overhead. Extensive experiments across three models and four benchmarks consistently validate our observation and the effectiveness of EB-Decode, achieving 3.53-18.76$\times$ higher throughput than the vanilla decoding method and up to 1.58$\times$ higher throughput over the strongest baseline, Fast-dLLM, with comparable accuracy.

1 Introduction

Diffusion LLMs offer parallel generation, but fixed blocks and confidence thresholds fail to account for heterogeneous token difficulty. EB-Decode learns variable-length blocks and early commitment, improving throughput over vanilla decoding with comparable accuracy.

1 Introduction

  • Fixed block selection and confidence thresholds waste parallelism on hard tokens and underuse it on easy ones, despite heterogeneous token difficulty.Existing approaches include fixed-block methods, heuristics, and learned sampling strategies that still rely on fixed block sizes.Source: §1 · pp. 1–2
  • EB-Decode groups similarly low-entropy, semantically close tokens for earlier joint decoding before they reach the confidence threshold.The framework is motivated by tokens with similar difficulty tending to cluster.Source: §1 · p. 2
  • LBS dynamically groups contiguous or non-contiguous tokens with similar difficulty and semantic coherence into variable-length blocks.It predicts block sizes on the fly rather than relying on heuristics.Source: §1 · p. 2
  • LPS uses a position-aware learnable sampler to identify converged tokens for early finalization within variable-length blocks.Both enablers are presented as components of the EB-Decode framework.Source: §1 · p. 2
  • 3.53–18.76× throughput improvement over the vanilla decoder, with comparable accuracy and ∼6.21% routing overhead, is reported across three dLLMs and four benchmarks.Combining EB-Decode with KV-cache pipelines delivers an additional 1.20× speedup.Source: §1 · p. 2

2 Related Works

Related work covers dLLM decoding, efficiency methods, and early prediction. EB-Decode is positioned as learning both token grouping and early commitment during dLLM decoding.

2 Related Works

  • dLLMs generate through iterative denoising over discrete tokens, and the cited literature includes LLaDA and Dream developments.The passage describes LLaDA as performing on par with LLaMA3 at billion-parameter scale.Source: §2 · p. 3
  • Efficiency work spans caching, distillation, sampling schedules, adaptive acceptance, dynamic blocks, and methods targeting test-time quality.EB-Decode is distinguished as learning to group tokens by difficulty and adapt sampling to finalize converged tokens early.Source: §2 · p. 3
  • Prior early-prediction research includes training-time early-bird tickets and inference-time early exits, while this work observes early-bird behavior during dLLM decoding.The passage notes observations across BERT, GCNs, LLMs, and diffusion models.Source: §2 · p. 3

3 Preliminaries of dLLMs

The preliminaries describe dLLM generation as forward masking followed by reverse denoising. Block-wise decoding predicts active-block tokens in parallel but commits them using confidence while retaining fixed block sizes.

3 Preliminaries of dLLMs

  • dLLMs model text generation as a diffusion process over discrete token sequences with forward masking and reverse denoising.The clean sequence is denoted x0 and has length L.Source: §3 · p. 3
  • The reverse process uses a bidirectional Transformer mask predictor to recover distributions over original tokens at masked positions simultaneously.The forward process independently replaces tokens with [MASK] at noise level t.Source: §3 · p. 3
  • Block-wise decoding processes fixed-size contiguous blocks from left to right, predicting active-block positions in parallel and committing based on confidence.Fixed blocks and static commit rules ignore variation in token difficulty along the sequence.Source: §3 · p. 3

4 The Observations of EB Phenomenon in dLLM Decoding

The section identifies staircase-like variation in token entropy across positions and early convergence across denoising steps. These patterns expose mismatches between token difficulty and fixed-block, confidence-threshold decoding.

4 The Observations of EB Phenomenon in dLLM Decoding

  • Entropy forms long, flat low-entropy plateaus interrupted by high-entropy spikes, rather than decreasing smoothly across token positions.The figure compares fixed-size blocks with non-contiguous learnable blocks; colored regions denote blocks and gray regions denote unmasked tokens.Source: §4 · p. 4, Figure 2 · p. 4
  • Non-contiguous grouping can defer hard tokens within coherent regions so surrounding easy context is decoded first and supplies later right-side information.Fixed boundaries can split semantically coherent regions and force hard tokens to be resolved before subsequent content can be generated.Source: §4 · p. 4
  • 13.4 and 2.3 steps are the respective mean lead times for vanilla and confidence-based methods after predictions become correct but before commitment.For lead time ≥1 step, the respective token fractions are 96.3% and 62.5%; for ≥10 steps, they are 60.5% and 3.5%.Source: §4 · pp. 4–5, Figure 3 · p. 4

5 The Proposed EB-Decode Framework

EB-Decode replaces fixed block sizes and confidence thresholds with learnable decisions while keeping the base dLLM frozen. Its LBS router selects variable-length blocks, and its LPS router learns which positions within them to finalize early.

5 The Proposed EB-Decode Framework

  • EB-Decode replaces fixed block size and confidence-threshold choices with learnable counterparts while keeping the base dLLM frozen.Both modules reuse the base model’s forward pass and add negligible training and inference overhead.Source: §5 · p. 5

5.1 Enabler 1: Learnable Block Size (LBS)

  • LBS selects variable-length, non-contiguous blocks from masked positions to align decoding with clustered token difficulty.The router uses normalized entropy and top-1 token IDs, combining uncertainty and token information for per-position inclusion decisions.Source: §5.1 · pp. 5–6
  • LBS training uses weighted cross-entropy to favor correct, low-loss positions and regularization to discourage selecting no tokens.The regularizer rewards larger total inclusion and balances the loss to stabilize block size.Source: §5.1 · p. 6
  • At inference, LBS removes EOS candidates, enforces a maximum gap between selected positions, and falls back to the first 32 masked positions when fewer than Lmin remain.These filters prevent premature EOS commitment, unresolved hard regions within a block, and overly sparse selections.Source: §5.1 · pp. 6–7

5.2 Enabler 2: Learnable Parallel Sampling (LPS)

  • LPS learns to finalize positions within a variable-length block before their confidence reaches the standard threshold.Its position-aware router uses relative intra-block position, confidence, normalized entropy, top-1/top-2 confidence gap, and block-wise mask ratio.Source: §5.2 · p. 7
  • LPS training replays the base model’s default decoding trajectory and weights incorrect early-commitment examples more heavily.At inference, positions commit when base confidence or the LPS score passes its threshold; the highest-confidence position is committed if neither does.Source: §5.2 · pp. 7–8

6 Experiments

Across three dLLMs and four benchmarks, EB-Decode improves throughput over vanilla decoding while maintaining comparable accuracy. Ablations and broader comparisons examine its components, sequence lengths, cache compatibility, transferability, and applicability to other decoding settings.

6.2 Comparison of EB-Decode with SOTA Baselines

  • 3.53–18.76× higher throughput than the vanilla decoder, with comparable accuracy, and up to 1.58× higher throughput than Fast-dLLM.The comparison covers three dLLMs and four benchmarks.Source: §6.2 · p. 9

6.3 Ablation Studies of EB-Decode

  • LBS alone raises HumanEval throughput from 23.70 to 68.83 TPS and accuracy from 43.90% to 45.73%.The broader ablation finds that filters refine the LBS-selected block, while learned block selection itself brings large gains; entropy-only input lowers accuracy by 1.2 points on average at similar TPS.Source: §6.2 · p. 9, §6.3 · pp. 9–10
  • LPS components—including a Transformer, intra-block position embeddings, and auxiliary uncertainty features—contribute to maximizing TPS while preserving task accuracy.The ablation builds on LBS-selected blocks, starting with an MLP that takes confidence as input.Source: Table 3 · p. 9, §6.3 · p. 10
  • EB-Decode achieves its highest throughput and comparable accuracy at generation lengths 256 and 512, with speedup growing at the longer length.The table compares accuracy and speedup at both lengths.Source: §6.3 · p. 10, Table 5 · p. 10
  • Prefix and dual KV caches further improve EB-Decode throughput while preserving accuracy; it remains 1.28×, 1.22×, and 1.38× faster than Fast-dLLM without cache, with prefix cache, and with dual cache, respectively.The cache strategies reuse keys and values at block boundaries.Source: §6.3 · p. 10, Table 6 · p. 10
  • The LBS and LPS routers add 3.80% and 2.41% of per-step backbone forward time, respectively, or 6.21% combined.The routers are measured on A100.Source: §6.3 · p. 10

6.4 Generality and Broader Comparisons

  • Transferability varies by setting: LBS transfers within the LLaDA series but not to Dream-v0, while LPS transfers across those models, generation lengths, and unseen code tasks.LBS is recommended at its training length; its token embeddings and positional embeddings are tied to the model vocabulary and generation length.Source: §6.4 · pp. 10–11, Table 7 · p. 11
  • On block-forward dLLMs, LBS improves accuracy in all eight tested configurations while throughput stays close to baseline; in a matched setting, EB-Decode runs about 2.4× faster than the faster AR model.For block-causal attention, LBS uses gmax = 1; the AR comparison uses batch size 1 on 4×A100 GPUs.Source: §6.4 · p. 11, Table 8 · p. 11, Table 9 · p. 12

7 Conclusion

The conclusion presents EB-Decode as a plug-in framework that replaces fixed block sizes and confidence thresholds with learnable routers while keeping the base dLLM frozen. The accompanying background and appendix passages describe dLLM decoding, its fixed-block baseline, and the LBS router architecture.

7 Conclusion

  • EB-Decode replaces fixed block size and confidence thresholds with lightweight learnable LBS and LPS routers, keeping the base dLLM frozen.Experiments across three open-source dLLMs and four benchmarks report up to 18.76× throughput improvement over vanilla decoding with comparable accuracy.Source: §7 · p. 12

A Detailed dLLM Backgrounds

  • The background traces masked dLLMs from discrete-diffusion objectives through later model extensions, including semi-AR methods that sample in parallel within autoregressive blocks.It describes D3PM, MDLM, MD4, SEDD, and subsequent masked-diffusion developments.Source: A Detailed dLLM Backgrounds · p. 17

B Block-wise Decoding for dLLMs

  • In the block-wise baseline, masked positions in a fixed-size active block are predicted in parallel, then selectively committed and refined until the block is complete.The sequence is denoised from t = 1 to t = 0; EB-Decode replaces the fixed block size and static commit rule with learnable counterparts.Source: B Block-wise Decoding for dLLMs · p. 17

C.1 Router Architecture

  • The LBS router is a 600k-parameter Transformer encoder that maps per-masked-position features to scalar inclusion logits.Its architecture uses entropy and token embeddings, fusion, a two-layer Transformer, and an output head.Source: §C.1 · pp. 17–18

C.2 Training Procedure

The LBS router is trained with the dLLM frozen, while the LPS router is a small position-aware Transformer trained from decoding traces. The supplied passages report stable training and trace supervision derived from completed sequences.

C.2 Training Procedure

  • The LBS router trains on dLLM-generated responses with the base dLLM frozen, using blockwise masking and updating only the router.Training is performed on 8×A100 GPUs.Source: §C.2 · p. 18
  • Both LBS-router losses drop sharply in the first epoch and change much less afterward, indicating stable training.The negative training loss reflects subtraction of the regularization term from weighted cross-entropy.Source: §C.2 · p. 18, Table 10 · p. 19

D.1 Router Architecture

  • The LPS router uses two pre-norm Transformer blocks with four attention heads and about 27k trainable parameters.Its final linear head produces a scalar logit for each position.Source: §D.1 · p. 19

D.2 Trace Generation

  • LPS training traces match inference settings, and correctness labels are derived by comparing intermediate predictions with the completed sequence.The replay procedure provides supervision without external annotation.Source: §D.2 · p. 19

D.3 Training Procedure

LPS training uses decoding traces, weighted binary cross-entropy, and validation-based threshold selection. At inference, EB-Decode combines learned block selection with parallel token commits and a fallback to prevent stalling.

D.3 Training Procedure

  • Valid trajectories are split 80/10/10 for training, validation, and testing, and the LPS threshold is selected for the largest recall under the false-positive-rate budget.The threshold search uses a 250-point grid from 0.8 to 0.9999.Source: §D.3 · pp. 19–20
  • Negative samples receive weight 2 because false-positive commits can irreversibly accept a wrong token, whereas false negatives only delay acceptance.The LPS threshold is tuned under a false-positive-rate budget of 10^-2.Source: §D.3 · pp. 19–20
  • LPS training and validation losses converge smoothly while validation recall rises, and the validation loss remains close to training loss.The authors interpret the close losses as indicating no overfitting.Source: §D.3 · p. 20

E EB-Decode Inference Procedure

  • Inference selects an active block using router scores and a max-gap filter, then commits accepted tokens in parallel; an empty commit set falls back to the highest-confidence masked token.The procedure also falls back to a fixed number of positions if the selected block is smaller than the minimum size.Source: E EB-Decode Inference Procedure · p. 21

F Experimental Setup

The experiments use shared decoding settings and report throughput per GPU at batch size 1. The setup also specifies router thresholds, hardware, and cache handling across methods.

F Experimental Setup

  • All methods use generation length 512, base block size 32, and confidence threshold τ=0.9; no method in Table 1 uses a KV cache.Fast-dLLM's official implementation recomputes the cache at each block start, while vanilla LLaDA maintains no cache.Source: F Experimental Setup · p. 21
  • LBS uses gmax=5 and minimum block sizes 8, 8, and 16 for LLaDA-8B-Instruct, Dream-v0-Instruct-7B, and LLaDA-1.5, respectively.The LPS acceptance thresholds are 0.89, 0.90, and 0.87 for those models, selected under a 1% ± 0.5 percentage-point false-positive-rate budget.Source: F Experimental Setup · p. 21
  • Throughput is generated tokens divided by end-to-end latency, measured per GPU with Hugging-Face Transformers at batch size 1.Runs use 4× NVIDIA H100 GPUs except specified experiments on A100 GPUs.Source: F Experimental Setup · pp. 21–22

G Positioning Against Prior Adaptive dLLM Decoding Methods

EB-Decode learns both where to decode and when to commit, unlike prior methods described as relying on fixed blocks or hand-designed heuristics and mostly static thresholds. The sequence tail also forms a persistent low-entropy region associated with EOS predictions.

G Positioning Against Prior Adaptive dLLM Decoding Methods

  • EB-Decode is the only compared method that learns both block formation and token commitment, supporting non-contiguous, variable-length blocks.The comparison covers where to decode and when to commit; the authors say hard tokens can be deferred for richer right-side context.Source: Table 12 · p. 22, G Positioning Against Prior Adaptive… · p. 22

H Entropy at the Sequence Tail

  • 96.30% of sampled positions among the sequence's last five have <|endoftext|> among the top-five predictions, forming a persistent low-entropy tail band.The table aggregates top predictions across decoding steps for LLaDA-8B-Instruct.Source: Table 13 · p. 22, H Entropy at the Sequence Tail · p. 22

I Limitations

The reported study is bounded to generation lengths of 256 and 512 tokens and throughput measurements at batch size 1. EB-Decode also adds deployment complexity through two routers and several hyperparameters.

Generation Length

  • Results cover generation lengths of 256 and 512 tokens; substantially longer-generation behavior remains for future study.The authors report a speedup trend from 256 to 512 tokens but call a more comprehensive study at substantially longer lengths a future extension.Source: I Limitations · p. 22

Serving Conditions

  • Throughput is measured at batch size 1; evaluation with continuous batching and system-level optimizations is left to future work.The authors note that this evaluation depends on dLLM serving infrastructure as well as the decoding algorithm.Source: I Limitations · p. 22

Method Complexity

  • Two routers and four hyperparameters make EB-Decode more involved to deploy and reuse than a single-threshold heuristic.The listed hyperparameters are gmax, Lmin, τ, and τψ; the authors identify reducing the number of moving parts as a worthwhile direction.Source: I Limitations · p. 23
Loading 2609.16450v1…