LLaDA2.1: Speeding Up Text Diffusion via Token Editing — part 1/2
Cleaned-text code points [0, 15574). Bibliography omitted here; source.md includes it.
# LLaDA2.1: Speeding Up Text Diffusion via Token Editing
###### Abstract
While LLaDA2.0 showcased the scaling potential of 100B-level block-diffusion models and their inherent parallelization, the delicate equilibrium between decoding speed and generation quality has remained an elusive frontier. Today, we unveil LLaDA2.1 , a paradigm shift designed to transcend this trade-off. By seamlessly weaving Token-to-Token (T2T) editing into the conventional Mask-to-Token (M2T) scheme, we introduce a joint, configurable threshold-decoding scheme. This structural innovation gives rise to two distinct personas: the Speedy Mode (S Mode) , which audaciously lowers the M2T threshold to bypass traditional constraints while relying on T2T to refine the output; and the Quality Mode (Q Mode) , which leans into conservative thresholds to secure superior benchmark performances with manageable efficiency degrade. Furthering this evolution, underpinned by an expansive context window, we implement the first large-scale Reinforcement Learning (RL) framework specifically tailored for dLLMs, anchored by specialized techniques for stable gradient estimation. This alignment not only sharpens reasoning precision but also elevates instruction-following fidelity, bridging the chasm between diffusion dynamics and complex human intent. We culminate this work by releasing LLaDA2.1-Mini (16B) and LLaDA2.1-Flash (100B) . Across 33 rigorous benchmarks, LLaDA2.1 delivers strong task performance and lightning-fast decoding speed. Despite its 100B volume, on coding tasks it attains an astounding 892 TPS on HumanEval+, 801 TPS on BigCodeBench, and 663 TPS on LiveCodeBench.
## 1 Introduction
Discrete diffusion Large Language Models (dLLMs) have emerged as a compelling alternative to autoregressive generation, offering the potential for non-monotonic reasoning and parallel decoding. However, the standard absorbing-state framework—which enforces a rigid, monotonic transition from [MASK] to fixed tokens—faces inherent limitations in fidelity. As highlighted by Kang et al. (2025) , the independent nature of parallel decoding often amplifies token-level inconsistencies. While recent studies have attempted to mitigate this via confidence-based remasking ( Wang et al., 2025b ) or by employing external guide models ( Lee et al., 2025 ) . To bridge the gap between efficient parallel generation and high-fidelity reasoning, we align with the direction of generalizing discrete diffusion beyond absorbing states ( Rütte et al., 2025 ) and propose a comprehensive framework for Editable State Evolution.
Unlike prior work such as Song et al. (2025) , we first design a novel Error-Correcting Editable decoding strategy, which introduces a dynamic paradigm controlled by dual probability thresholds. This paradigm encompasses two types of operations: direct decoding from mask to token, and editing from one token to another. This strategy enables the model to directly refine its own outputs during the generation process, thereby effectively addressing the local inconsistencies commonly encountered in parallel decoding. To cultivate this editing capability, our CPT and SFT phases expose the model to both masked positions and stochastic noise, incentivizing it to not only generate new content but also identify and rectify existing errors.
Crucially, this architecture transforms the rigid trade-off between latency and fidelity into a flexible, user-configurable continuum. By allowing the model to retroactively correct errors, we can aggressively lower the confidence threshold for the initial Mask-to-Token (M2T) phase without collapsing the generation quality. This insight gives rise to two distinct operating personas: a Speedy Mode (S Mode) , which prioritizes high-throughput generation by accepting lower-confidence tokens and relying on subsequent Token-to-Token (T2T) passes for rectification; and a Quality Mode (Q Mode) , which adheres to conservative thresholds to maximize reasoning rigor. This duality demonstrates that editability is not merely a mechanism for error repair, but a fundamental lever for accelerating parallel decoding.
To further elevate the model’s capabilities, we integrate a Reinforcement Learning (RL) stage. While recent works such as SPG ( Wang et al., 2025a ) , TraceRL ( Wang et al., 2025c ) and ESPO ( Ou et al., 2025 ) have demonstrated the potential of RL in improving dLLMs, applying policy gradients to block-autoregressive models remains challenging due to the intractability of sequence log-likelihoods. We circumvent this by adopting an ELBO-based Block-level Policy Optimization (EBPO) framework tailored for our editable setting.
Notice that LLaDA2.1 extends its previous version (LLaDA2.0) by prioritizing decoding versatility over mere parameter scaling or benchmark peaking. By keeping the model size constant and minimal change of training data, we prove that our novel editing scheme enables lightning-fast execution with minimal overhead. This work serves as a proof-of-concept for a new dLLM paradigm that balances high-quality generation with extreme operational efficiency.
## 2 Configurable Decoding Scheme
During LLM decoding, Exposure Bias —where errors compound as the model conditions on its own imperfect predictions—is inevitable. This phenomenon is particularly severe in dLLMs due to their parallel generation nature. We observe that once such decoding errors occur, dLLMs tend to become increasingly conservative in subsequent steps, significantly slowing down the generation process. In contrast, autoregressive models exhibit lower exposure bias and can self-correct through extended chain-of-thought reasoning. To address this challenge, we introduce an editing operation into the decoding process, enabling the model to retrospectively correct errors introduced during parallel generation, thereby achieving a much better balance between generation speed and quality.
Specifically, we extend standard discrete diffusion to support it. Unlike conventional absorbing-state models that enforce a rigid monotonic transition from [MASK] to fixed tokens, our framework introduces a dynamic “Draft-and-Edit” paradigm controlled by dual probability thresholds. We formalize the state evolution by defining two active update sets at timestep t t : the Unmasking Set Γ t \Gamma_{t} and the Editing Set Δ t \Delta_{t} .
We formalize the state evolution by defining two active update sets at timestep t t : the Unmasking Set Γ t \Gamma_{t} and the Editing Set Δ t \Delta_{t} . Let v t i = arg max v p θ ( v | 𝒙 t ) v_{t}^{i}=\arg\max_{v}p_{\theta}(v|\bm{x}_{t}) be the top-candidate. The update indices are identified as: Γ t \displaystyle\Gamma_{t} = { i ∣ x t i = [MASK] and p θ ( v t i | 𝒙 t ) > τ mask } , \displaystyle=\left\{i\mid x_{t}^{i}=\text{[MASK]}\text{ and }p_{\theta}(v_{t}^{i}|\bm{x}_{t})>\tau_{\text{mask}}\right\}, (1) Δ t \displaystyle\Delta_{t} = { i ∣ x t i ≠ v t i and p θ ( v t i | 𝒙 t ) > τ edit } , \displaystyle=\left\{i\mid x_{t}^{i}\neq v_{t}^{i}\text{ and }p_{\theta}(v_{t}^{i}|\bm{x}_{t})>\tau_{\text{edit}}\right\}, (2) with τ mask , τ edit ∈ [ 0 , 1 ] \tau_{\text{mask}},\tau_{\text{edit}}\in[0,1] being the confidence thresholds configuring the decoding dynamics. The transition operator then applies the updates strictly on the union of these sets: x t − 1 i = { v t i if i ∈ Γ t ∪ Δ t , x t i otherwise . x_{t-1}^{i}=\begin{cases}v_{t}^{i}&\text{if }i\in\Gamma_{t}\cup\Delta_{t},\\ x_{t}^{i}&\text{otherwise}.\end{cases} (3)
## 3 Training Paradigm
### 3.1 Training Alignment for “Draft-and-Edit”
To align the model with the “Draft-and-Edit” inference paradigm and mitigate the Exposure Bias inherent in standard mask-based training, we employ a unified Mixture of M2T and T2T objective. This objective is applied throughout both the Continual Pre-Training (CPT) and Supervised Finetuning (SFT) stages.
This dual-stream training objective enables the model to develop two complementary capabilities fundamental to our framework:
• Drafting Stream (Mask-to-Token): The model learns to predict the correct token at each masked position to generate initial content, establishing the foundational drafting capability.
• Editing Stream (Token-to-Token): The model learns to recover original tokens from random noise perturbations (rectifying errors), equipping it with the ability to identify and rewrite artifacts.
By consistently applying this dual-stream supervision from CPT through SFT, we ensure that LLaDA2.1 is fundamentally conditioned to function as both a fast drafter and a precise editor within a single parameter space. Additionally, we employ a Multi-turn Forward (MTF) data augmentation technique, by exposing the model to a wider variety of editing scenarios, enhance the model’s editing capabilities.
### 3.2 Reinforcement Learning Training
The application of policy gradient methods to diffusion models faces a fundamental hurdle: the intractability of the sequence-level log-likelihood, log π θ ( 𝒙 ) \log\pi_{\theta}(\bm{x}) , which is essential for computing policy updates. While prior works have explored various approximations, they have historically struggled with high variance and prohibitive computational costs, limiting RL to small-scale experiments ( Wang et al., 2025c , Ou et al., 2025 , Wang et al., 2025a ) . We overcome this bottleneck by synthesizing ELBO-based Block-level Policy Optimization (EBPO) with robust infrastructure optimizations. By utilizing the Evidence Lower Bound (ELBO) as a principled proxy for exact likelihood and implementing Vectorized Likelihood Estimation ( Arriola et al., 2025 ) to parallelize bound computation, we achieve orders-of-magnitude acceleration. This integration allows us to scale dLLMs RL to unprecedented context lengths and training magnitudes, establishing a stable and efficient pipeline for post-training.
Formally, we maximize a clipped surrogate objective, where the advantage is weighted by the probability ratio ρ \rho : 𝒥 EBPO ( θ ) = 𝔼 𝒙 , 𝒚 ∼ π θ old [ min ( ρ ( 𝒚 | 𝒙 ) A ^ , clip ( ρ ( 𝒚 | 𝒙 ) , 1 − ϵ low , 1 + ϵ high ) A ^ ) ] , \mathcal{J}_{\text{EBPO}}(\theta)=\mathbb{E}_{\bm{x},\bm{y}\sim\pi_{\theta_{\text{old}}}}\left[\min\left(\rho(\bm{y}|\bm{x})\hat{A},\text{clip}(\rho(\bm{y}|\bm{x}),1-\epsilon_{\text{low}},1+\epsilon_{\text{high}})\hat{A}\right)\right], (4) where A ^ \hat{A} is an estimator of the advantage function at timestep t t , quantifying the relative improvement of the chosen action over the average expectation under the current policy. For a set of discretized timesteps { t n } n = 1 N \{t_{n}\}_{n=1}^{N} and weights { w n } \{w_{n}\} , we construct a composite input 𝒛 n = 𝒚 t n ⊕ 𝒚 0 \bm{z}_{n}=\bm{y}_{t_{n}}\oplus\bm{y}_{0} to compute all block-conditional probabilities in parallel: log ρ ( 𝒚 | 𝒙 ) ≈ ∑ n = 1 N w n ∑ b = 1 B ( log p θ ( 𝒚 b ∣ 𝒛 n , 𝒙 ; ℳ ) − log p θ old ( 𝒚 b ∣ 𝒛 n , 𝒙 ; ℳ ) ) . \log\rho(\bm{y}|\bm{x})\approx\sum_{n=1}^{N}w_{n}\sum_{b=1}^{B}\left(\log p_{\theta}(\bm{y}^{b}\mid\bm{z}_{n},\bm{x};\mathcal{M})-\log p_{\theta_{\text{old}}}(\bm{y}^{b}\mid\bm{z}_{n},\bm{x};\mathcal{M})\right). (5) Here, ℳ \mathcal{M} denotes a Block-Causal Mask ensuring the b b -th block attends only to valid history. By aggregating block-level contributions ( ∑ b = 1 B \sum_{b=1}^{B} ) within a single forward pass per timestep n n , we establish a computationally tractable pipeline for scaling reinforcement learning to long-context diffusion generation.
## 4 Infrastructure
### 4.1 Training Infrastructure
#### Continued Pre-Training and Supervised Fine-Tuning
For both continued pre-training (CPT) and supervised fine-tuning (SFT), we adopt the same training infrastructure as LLaDA2.0 ( Bie et al., 2025 ) , leveraging dFactory ( InclusionAI, 2025 ) , which provides efficient training recipes specifically designed for dLLMs, except that we introduce a dedicated optimized implementation for the multi-turn forward (MTF) stage.
#### RL Training
To enable effective policy optimization for dLLMs, we extend the AReaL framework ( Fu et al., 2025 , Mei et al., 2025 ) by developing specialized likelihood estimation and advantage estimation protocols that leverage diffusion sampling, explicitly supporting both T2T and M2T modes. This workflow is powered by ASystem ( Ling Team and others, 2025 ) for distributed orchestration and utilizes a customized version of SGLang ( Ant Group Team and SGLang Team, ) as the dedicated rollout engine.
### 4.2 Inference Infrastructure
We use a customized version of SGLang ( Ant Group Team and SGLang Team, ) for inference. To further accelerate the inference speed, we integrate Alpha-MoE ( Aleph-Alpha, ) , a MoE megakernel that combines the two FusedMoE computations into one kernel, and adopt per-block FP8 quantization to balance the inference speed and model accuracy. To accelerate inference on long-context sequences, we adopt block-wise causal masked attention, allowing the KV cache for the entire long context to be computed in a single forward pass. We further enable radix caching and batching support for block diffusion LLMs in SGLang.
### 4.3 Decoding Algorithm at Inference
In the inference stage, we adopt a decoding algorithm that combines Threshold Decoding ( Ma et al., 2025 ) with an explicit editing mechanism. In the basic setting, decoding and editing are performed within a single block: tokens are generated under a threshold-based constraint, and local edits are applied to revise intermediate outputs before the block is finalized.
Beyond single-block editing, we further introduce a Multiple Block Editing (MBE) mechanism. MBE allows the model to revisit and revise previously generated blocks based on the content of newly decoded blocks.
## 5 Evaluation
To comprehensively evaluate the quality of instruction-tuned models, we employ a diverse suite of benchmarks categorized into five dimensions:
• Knowledge : MMLU-Pro ( Wang et al., 2024 ) , GPQA-Diamond ( Rein et al., 2024 ) , C-Eval ( Huang et al., 2023 ) , PHYBench ( Qiu et al., 2025 ) , TriviaQA ( Joshi et al., 2017 )
• Reasoning : SQuAD 2.0 ( Rajpurkar et al., 2018 ) , DROP ( Dua et al., 2019 ) , KOR-Bench ( Ma et al., 2024 ) , HellaSwag ( Zellers et al., 2019 ) , BIG-Bench Hard ( Suzgun et al., 2023 ) , BIG-Bench Extra Hard ( Kazemi et al., 2025 ) , MuSR ( Sprague et al., 2023 ) , ZebraLogic ( Lin et al., 2025 ) , PrOntoQA ( Saparov and He, 2022 ) , PIQA ( Bisk et al., 2020 ) , OCNLI ( Hu et al., 2020 ) , BIG-Bench Hard-CN ( Opencompass Team, 2023 )
• Coding : CRUXEval ( Gu et al., 2024 ) , MultiPL-E ( Cassano et al., 2023 ) , BigCodeBench ( Zhuo et al., 2024 ) , LiveCodeBench ( Jain et al., 2024 ) , Spider ( Yu et al., 2018 ) , BIRD ( Li et al., 2023 ) , HumanEval+ ( Liu et al., 2023 ) , MBPP+ ( Liu et al., 2023 )
• Math : OlympiadBench ( He et al., 2024 ) , AIME 2025 ( AIME, 2025 ) , Omni-MATH ( Gao et al., 2024 ) , GSM-Plus ( Li et al., 2024 ) , CMATH ( Wei et al., 2023 )
• Agent & Alignment : BFCL ( Patil et al., 2025 ) , IFEval ( Zhou et al., 2023 ) , Nexus Function Calling Benchmark ( Nexusflow.ai Team, 2023 )
We report the comparative scores and TPF (tokens per forward) of LLaDA2.1-flash and LLaDA2.1-mini against other models in Tables 1 and 2 , respectively. From the results, we observe that LLaDA2.1’s scores under S Mode decrease compared to LLaDA2.0, but a substantial improvement in TPF is achieved. While under Q Mode , LLaDA2.1 surpasses the results of LLaDA2.0 on both mini and flash model.