> Assessing masked diffusion language models viability for multi‑token prediction, diagnosing their efficiency and learning challenges, and surveying emerging solutions within the broader multi‑token prediction landscape. We seek to consolidate our recent reflections on diffusion language models (DLLMs), with particular emphasis on the mainstream semi‑autoregressive, mask‑based variants. Whether DLLMs can ultimately replace autoregressive (AR) models as the default backbone for large language models remains an open question; nevertheless, they offer a promising alternative for multi‑token prediction tasks. Present‑day open-source DLLMs confront two principal challenges. First, the generation efficiency of pure DLLMs is still under debate: their architectures do not naturally accommodate key‑value caching, and overly aggressive reductions in the number of denoising function evaluations (NFEs) can markedly impair generation quality. Second, because a mask‑based DLLM can be regarded as any‑order AR model, it faces a higher learning complexity on textual data than a pure AR counterpart. This note reviews these challenges, surveys current exploratory solutions, and discusses how DLLMs intersect with the broader paradigm of multi‑token prediction. ## 1 Autoregressive Models In classical autoregressive (AR) language models, training is typically performed with **next-token prediction** under **teacher forcing**. The learning objective is $$ -\log p_{\boldsymbol{\theta}}(\boldsymbol{x})=-\sum_{i=1}^n \log p_{\boldsymbol{\theta}}\left(\boldsymbol{x}^i \mid \boldsymbol{x}^{t_{1}>\dots>t_{K}=0$: | 1. Predict the masked tokens given the current $x_{t_{k}}$. 2. Retain high-confidence predictions while re-masking low-confidence positions (the “remask” strategy) to preserve stochasticity. 3. Proceed to the next time step $t_{k+1}$ After several iterations the process yields the final text $x_0$. In **LLaDA**, the confidence filter keeps a fixed number of tokens in each denoising step: for a target length $N$ and $T$ denoising steps, exactly $\max\{1,\,N/T\}$tokens are preserved per step. | ### 2.4 Semi-Autoregressive MDLLM **Semi-AR-MDLLM** integrates conventional autoregressive (AR) decoding with a Mask-Based Diffusion Language Model (MDLLM). Its logic is deliberately straightforward： 1. **Blocking/Chunking** Partition the target sequence into equal-length segments, denoted _Block 0, Block 1, …, Block n_ ; 2. **Inter-block Autoregression** Generate the blocks left-to-right—first _Block 0_, then _Block 1_, and so on—thereby preserving causality across blocks； 3. **Intra-block Diffusion** Within each block, tokens are denoised through MDLLM’s unmasking mechanism ; This design yields a trade-off between pure AR and full MDLLM decoding: - **Block length $= 1$ $\rightarrow$ pure AR decoding**； - **Block length $=$ sequence length** $\rightarrow$ **single-shot MDLLM**； The figure blow schematically illustrates the semi-autoregressive sampling schedule. At the initial step ($t=1$) all tokens except the prompt are masked. The model then successively reconstructs _Block 0 -> Block 1 -> Block 2 -> ⋯_, reaching a fully realized sequence at $t=0$. ![image.png](/images/dllm_images/image1.png) - It’s worth noting that In LLada’s implementation, even under semi-AR generation the entire sequence is fed to the model at **every** forward pass. When generating _Block 0_, for example, the input is _Prompt + Block 0 + Block 1 + Block 2 + ⋯_, rather than _Prompt + Block 0_ alone. We conjecture that this choice ensures consistency with the supervised fine-tuning (SFT) regime, which conditions on the full sequence. - If the training phase employed _Prompt + Block 0 + … + Block i_ for predicting tokens of the _Block i + 1_, then inference could naturally exploit **KV-caching / semi KV caching**, as only the past blocks would need to be re-introduced. In the latter Section, we will introduce BD3LM which utilized this training paradigm. ### 2.5 Any-Order Autoregressive (**AO-AR)** In **RADD** [4], the authors give a formal proof that the optimization objective of a **Mask-Based Diffusion Language Model** (MDLLM) is equivalent to that of an **arbitrary-order autoregressive** (AO-AR) model—exemplified by **$\sigma$-GPT** [14]. Readers interested in the full derivation are referred to the original paper [4]. Specifically, the AO-AR loss can be expressed as $$ -\log p_{\theta}(\boldsymbol{x}) = -\log \mathbb{E}_{\sigma \sim \mathcal{U}(S_n)} p_{\theta}(\boldsymbol{x} \mid \boldsymbol{\sigma}) \le \mathbb{E}_{\sigma \sim \mathcal{U}(S_n)} \left[ - \sum_{i=1}^{n} \log p_{\theta}(\boldsymbol{x}_{\sigma_i} \mid \boldsymbol{x}_{\sigma_{ \frac{n}{n+1}$. - **Practical implementation strategy** - **Fixed-threshold control.** In practice, **Fast-DLLM** does not enforce the theorem’s exact condition. Instead, it regulates the degree of parallel decoding in each denoising step with a fixed probability threshold. - **Concrete procedure (threshold = 0.9).** - If a masked position attains a predicted probability **> 0.9**, its token is accepted immediately—regardless of how many unmasked tokens have already been accepted in the current step. - Otherwise, the algorithm checks whether **any** masked position exceeds the 0.9 threshold. If none does, it accepts **the token with the highest predicted probability**, thereby guaranteeing that at least one token is produced in every step. - **Empirical performance.** This heuristic substantially reduces computational overhead and performs well in practice. Although it deviates from the theorem’s strictly greedy criterion, system-level experiments indicate that the approach is effective. - **Strict implementation under the theorem.** - We reproduced the theorem’s exact procedure in code and carried out a small-scale test; in practice it **does work.** - Concrete algorithmic steps: - **Collect greedy-unmasking probabilities.** For every masked region (block), perform one round of greedy unmasking at all masked positions and record the probability assigned to the selected token at each position. Example. For the sequence “今天吃 [mask1][mask2]”, suppose the model’s distribution for **[mask1]** is `{"什": 0.7, "鱼": 0.2, "其他": 0.1}`. Greedy unmasking chooses the token “什” with probability 0.7. - **Sort in descending order.** Gather the selected tokens across all masked positions and sort them by their probabilities. Let $a = [0.90,\,0.85,\,0.81,\,0.70,\,0.40,\,0.10,\,0.05,\,0.01]$, index $=[2, 3, 0, 1, 4, 6, 7, 5]$ - **Construct a threshold sequence and filter.** By Theorem 1, define $$ b = \Bigl[0,\,\tfrac23,\,\tfrac34,\,\tfrac45,\,\tfrac56,\,\tfrac67,\,\tfrac78,\,\tfrac89\Bigr] $$ (The first entry is set to 0 because each denoising step must accept at least one prediction, so the highest-probability token is always retained.) Identify the largest index $k$ such that $a_k>b_k$. In this example $k=2$. - **Retain high-confidence tokens.** Keep only the tokens whose probabilities are 0.90, 0.85, 0.81 (indices 2, 3, 0). All other tokens produced by unmasking are **remasked.** ## 5 Multi-Token Prediction The central idea of **Multi-Token Prediction** is to emit several future tokens per decoding step while incurring negligible additional compute. Current implementations (the list is not exhaustive) fall into two broad categories: 1. **Multi-head parallel prediction** [11]. Additional prediction heads ($head_1 … head_k$) simultaneously generate $x^{t+1},\dots,x^{t+k}$. A high-confidence teacher model then verifies and filters the candidates token by token — i.e., speculative decoding. A representative example is **Medusa** [12]. 2. **Register-token schemes**. Register tokens are periodically inserted into the input; their hidden states are used to predict several subsequent regular tokens. A recent instance is **MuToR** [7]. (MuToR is aimed not at higher throughput but at **improved performance**; recent work [11] suggests that MTP training, by injecting forecasts of future tokens, provides stronger planning signals. These issues, however, lie outside the focus of the present note.) ![image.png](/images/dllm_images/image9.png) For today’s mainstream **Semi-AR MDLLM** variants, their behavior is in many ways reminiscent of **register-token** MTP schemes: - The right-hand panel (excerpted from MuToR) can be read as predicting, at each step, **one regular token and one register token**. Both tokens attend only to themselves and the preceding regular tokens. The register token forecasts a future token that is offset by a pre-defined shift $d$. - Suppose we ignore Fast-DLLM’s threshold-driven parallel decoding and adopt the simplest policy—generating a fixed $C$ tokens per denoising step—and likewise set the block size to $C$. When processing block $k$, we can treat every **[MASK]** position as a register token and assign to each a distinct random shift $d\in[1,C]$. Under this construction, Semi-AR decoding reduces exactly to the second type of MTP outlined above. - In a conventional AR model, position $i$ predicts token $x^{i+1}$; in an MDLLM, position $i$ typically predicts $x^i$. This is not a hard rule, however: for example, in [13] the authors align MDLLM positions to match AR decoding. ## Contributions - [Yezhen Wang](https://scholar.google.com/citations?user=g-VEnLEAAAAJ&hl=zh-CN) - [Bo Li](https://brianboli.com/) ## Reference [1] Large Language Diffusion Models [2] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution [3] Any-Order GPT as Masked Diffusion Model: Decoupling Formulation and Architecture [4] Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data [5] Score-Based Continuous-Time Discrete Diffusion Models [6] Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding [7] Multi-Token Prediction Needs Registers [8] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models [9] Structured Denoising Diffusion Models in Discrete State-Spaces [10] Simple and Effective Masked Diffusion Language Models [11] Better & Faster Large Language Models via Multi-token Prediction [12] Medusa: Simple llm inference acceleration framework with multiple decoding heads [13] Scaling Diffusion Language Models via Adaptation from Autoregressive Models [14] $\sigma$-GPTs: A New Approach to Autoregressive Models --- ## Citations ```bibtex @misc{wang2025, author = {Wang, Yezhen and Li, Bo}, title = {Assessing Diffusion LM in the view of Multi-token Prediction}, year = {2025}, howpublished = {LMMS Lab Notes}, url = {https://www.lmms-lab.com/notes/dllm/}, note = {Accessed on 2025-07-23} } ```

Assessing Diffusion LM in the view of Multi-token Prediction

The Dream of the Digital Tide