Source-linked AI summary
LLaDA2.1: Speeding Up Text Diffusion via Token Editing
Tiwei Bie, Maosong Cao, Xiang Cao, Bingsen Chen, Fuyuan Chen, Kun Chen, Lun Du, Daozhuo Feng, Haibo Feng, Mingliang Gong, Zhuocheng Gong, Yanmei Gu, Jian Guan, Kaiyuan Guan, Hongliang He, Zenan Huang, Juyong Jiang, Zhonghui Jiang, Zhenzhong Lan, Chengxi Li, Jianguo Li, Zehuan Li, Huabin Liu, Lin Liu, Guoshan Lu, Yuan Lu, Yuxin Ma, Xingyu Mou, Zhenxuan Pan, Kaida Qiu, Yuji Ren, Jianfeng Tan, Yiding Tian, Zian Wang, Lanning Wei, Tao Wu, Yipeng Xing, Wentao Ye, Liangyu Zha, Tianze Zhang, Xiaolu Zhang, Junbo Zhao, Da Zheng, Hao Zhong, Wanli Zhong, Jun Zhou, Junlin Zhou, Liwang Zhu, Muzhi Zhu, Yihong Zhuang
TL;DR
LLaDA2.1 targets the speed–quality tension and fidelity problems of parallel dLLM decoding. It combines editable threshold decoding with block-level RL, achieving strong benchmark performance and substantial speed advantages while retaining mode-dependent trade-offs.
Problem
Parallel dLLM decoding can amplify token inconsistencies and exposure-bias errors, creating a challenge for balancing generation fidelity with efficiency.
Method
The paper combines M2T drafting, T2T error-correcting edits, configurable thresholds, and ELBO-based block-level policy optimization for dLLMs.
Results
Across evaluations, S Mode substantially improves throughput with only negligible quality sacrifice, while Q Mode surpasses LLaDA2.0 on mini and flash models.
Takeaways & Limitations
Editable decoding makes speed and quality configurable, with S Mode especially effective in code and Q Mode preferable for some general-chat cases.
Takeaways & Limitations
Threshold parameters require domain-specific adjustment, and S Mode can produce undesirable outputs in some general-chat cases.
Abstract
from arXiv · showhide
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
LLaDA2.1 addresses fidelity problems in parallel dLLM decoding by combining Mask-to-Token generation with Token-to-Token error correction. Configurable thresholds support either faster decoding or more conservative quality, complemented by a specialized RL stage.
- Parallel dLLM decoding can amplify token inconsistencies because standard absorbing-state transitions move rigidly from [MASK] to fixed tokens.
- The Error-Correcting Editable strategy combines direct mask-to-token decoding with token-to-token editing so the model can refine outputs during generation.CPT and SFT expose the model to masked positions and stochastic noise to support both generation and error correction.
- Lower M2T thresholds can increase throughput while later T2T passes rectify errors, whereas conservative thresholds prioritize reasoning rigor.These settings define Speedy Mode and Quality Mode as user-configurable operating personas.
- EBPO circumvents the difficulty of sequence log-likelihoods in block-autoregressive diffusion models to provide a tailored RL framework.
- With model size held constant and minimal training-data changes, the editing scheme targets fast execution with limited overhead.
2 Configurable Decoding Scheme
LLaDA2.1 extends discrete diffusion with a Draft-and-Edit decoding process that corrects errors introduced by parallel generation. Dual confidence thresholds determine which masked or existing tokens are updated.
- Exposure bias compounds parallel-decoding errors, making dLLMs increasingly conservative and slowing later generation steps.Retrospective editing is introduced to correct errors and improve the speed–quality balance.
- The Draft-and-Edit paradigm replaces rigid monotonic mask transitions with dynamic unmasking and editing operations.The framework defines an Unmasking Set Γ_t and an Editing Set Δ_t at each timestep.
- The decoding dynamics use dual probability thresholds to select update indices for masked positions and token replacements.
- Updates are applied through the union of the selected unmasking and editing sets.
3 Training Paradigm
LLaDA2.1 trains a single model to draft from masks and edit noisy tokens, then uses stable block-level policy optimization for scalable diffusion-model RL. The resulting pipeline supports long-context post-training.
- A unified mixture of M2T and T2T objectives is applied during both continual pre-training and supervised fine-tuning.
- The drafting stream predicts tokens at masked positions, while the editing stream recovers original tokens from random noise.
- Dual-stream supervision conditions one parameter space to act as both a fast drafter and a precise editor.Multi-turn Forward augmentation broadens the editing scenarios seen during training.
- The RL method addresses intractable sequence-level log-likelihoods through an ELBO-based block-level policy-optimization framework.
- The integration scales dLLM reinforcement learning to unprecedented context lengths and training magnitudes with a stable, efficient post-training pipeline.
- A clipped surrogate objective weights policy updates by an advantage estimate and probability ratio.The advantage estimates the chosen action’s relative improvement over the current policy’s average expectation.
- Composite inputs and block-causal masking aggregate block-level probabilities within one forward pass per timestep.This makes long-context diffusion RL computationally tractable.
4 Infrastructure
The infrastructure combines established dLLM training and distributed RL systems with specialized inference optimizations. Decoding supports local and cross-block editing alongside threshold-based generation.
- Training Infrastructure: CPT and SFT reuse LLaDA2.0’s dFactory infrastructure, with an optimized implementation for multi-turn forward training.
- RL Training: RL training extends AReaL with likelihood and advantage estimation protocols supporting both T2T and M2T modes.ASystem handles distributed orchestration and customized SGLang performs rollouts.
- Inference: Inference combines customized SGLang, Alpha-MoE, per-block FP8 quantization, and block-wise causal attention to improve speed and accuracy trade-offs.
- Inference: Single-block editing generates threshold-constrained tokens and revises intermediate outputs before finalizing each block.
- Inference: Multiple Block Editing revisits earlier blocks using content from newly decoded blocks.
5 Evaluation
LLaDA2.1 is evaluated across diverse knowledge, reasoning, coding, math, and agent/alignment benchmarks, with results emphasizing configurable speed–quality trade-offs. S Mode substantially improves throughput, while Q Mode improves over LLaDA2.0 and Multi-Block Editing further improves performance with modest throughput cost.
- Benchmark coverage: The evaluation spans five dimensions: knowledge, reasoning, coding, math, and agent & alignment.The benchmark suite includes nine knowledge and reasoning categories, eight coding benchmarks, five math benchmarks, and three agent/alignment benchmarks.
- Quality and efficiency: S Mode improves TPF substantially but lowers comparative scores relative to LLaDA2.0.Under Q Mode, LLaDA2.1 surpasses LLaDA2.0 on both Mini and Flash models.
- Throughput: 891.74 TPS is LLaDA2.1-Flash’s peak throughput on HumanEval+, while LLaDA2.1-Mini reaches 1586.93 TPS after quantization.Speed varies by domain, with the highest throughput in coding and the lowest in instruction following.
- Editing enhancement: Multi-Block Editing consistently improves benchmark performance for Flash and Mini variants at the cost of a modest throughput reduction.The gains are particularly evident on reasoning and coding tasks, where iterative cross-block refinement improves global consistency.
- Throughput: LLaDA2.1 S Mode achieves dramatically faster inference while sacrificing only a negligible amount of output quality across five benchmark domains.Figure 3 compares LLaDA2.1 variants with LLaDA2.0, Ling, and Qwen3 across the Mini and Flash series.
6 Outlook and Limitation
LLaDA2.1 improves inference speed through editable decoding, but speed and accuracy remain domain-dependent and require threshold tuning. The model remains experimental, with unresolved issues around error correction and initial-draft quality.
- Speed-accuracy tradeoffs persist across domains, requiring threshold adjustment to balance inference speed and accuracy.S Mode performs well on structured code and math with little accuracy loss, whereas some general-chat settings can produce undesirable output.
- The evaluation tables organize benchmark scores, throughput, relative score changes, and Multi-Block Editing comparisons across model variants.Table 1 and Table 2 report benchmark scores with tokens per forward for diffusion models, while Table 3 and Table 4 expose throughput and editing tradeoffs.
- High parallelism introduces more errors than autoregressive decoding, which can slow later reasoning.The paper reports that editing can correct decoding errors and help maintain higher inference speeds.
- Aggressively lowering τmask can produce rough drafts and n-gram repetition artifacts, although self-correction partially alleviates them.Balancing drafting speed with initial-structure quality remains an operational frontier.
- LLaDA2.1’s editing feature lowers the dLLM decoding threshold through cumulative error correction and yields substantial inference-speed benefits.The conclusion also states that many issues remain unresolved.