Organizations: Dots Studio, Xiaohongshu · Polyadic Research Group · Peking University · SAGA Lab, The Chinese University of Hong Kong, Shenzhen · The University of California, Los Angeles
Continuous diffusion language models (CDLMs) offer an alternative approach to parallel language generation; however, most existing studies directly use transformer architectures from either causal large language models (LLMs) or Diffusion Transformers (DiTs) for image generation, which were not optimized for CDLMs. In this paper, we conduct a systematic study on Transformer architectures for CDLMs and introduce Language Diffusion Transformer (LangDiT), combining feature and query-key normalization, wider noise conditioning, likelihood-shortcut gradient routing and calibration, and confidence-gated training self-conditioning. LangDiT significantly reduces generation perplexity in matched LangFlow comparisons, establishing a new state of the art in continuous language generation across different model formulations. Using LangFlow as a representative CDLM, LangDiT as a plug-and-play backbone leads to a new state-of-the-art model with a generation perplexity (GenPPL) of 21.72 using 1024 NFEs and 44.72 GenPPL using 128 NFEs on the unconditional generation of 1024-token sequences on the OpenWebText dataset. When adapted to Plaid-style CDLM formulations, LangDiT results in a new leading GenPPL of 11.33 using 1024 denoising NFEs, and 27.51 GenPPL using 128 denoising NFEs. These gains are achieved while retaining each model's underlying representation, training objective, and inference sampler, isolating transformer architecture as a fundamental source of improvement. Our study identifies architectural design as an important source of improvement in continuous diffusion language models.
Figures & tables
Figure 1: Architectural ablations during training on OWT-1024. GenPPL (left) and mean per-sample unigram entropy (right) at 128 NFEs, using 1,024 samples and EMA weights. R1 (baseline) and R7 (full model) are shown through 180k updates. For clarity, trajectories for R2–R6 are shown only through 50k updates. Points are unsmoothed logged measurements; endpoints are labeled at 180k. Legend additions are cumulative, with wider conditioning preceding shortcut detachment (Table 4 ). Lower GenPPL is accompanied by lower unigram entropy; checkpoint steps and sample counts are stated separately for each comparison.
Figure 2: Generation results on OWT ( L=1024 ). (a) Baseline and LangDiT at 128/1,024 denoising steps using checkpoints within 250k updates. The LangFlow comparisons use 250k checkpoints at 128 NFE and 180k checkpoints at 1,024 NFE. RePlaID compares 40k versus 150k at 128 steps and 220k versus 250k at 1,024 steps, so these pairs do not match training progress; the latter also differs in planned budget. (b) Extended-training checkpoints: LangFlow 660k and RePlaID 520k from 1M-update plans. RePlaID requires one final denoiser call beyond the displayed step count. Bars show GenPPL; bracketed values show mean per-sample unigram entropy. All results use 1,024 samples and EMA weights; † denotes our implementation (Table 3 ).
Component
Design principle
Implementation
Backbone
Separate feature-scale control from noise-dependent modulation.
Backbone and Q/K RMSNorm; wider noise conditioning.
Vocabulary head
Control the shortcut’s forward contribution and direct codebook gradients separately.
Shortcut detachment; bounded learned calibration.
Self-conditioning
Distinguish clean-state estimation from auxiliary feedback during training.
Confidence-gated bootstrap inputs; unweighted posterior-mean estimates at inference.
Table 1: Architectural design principles and their implementation in LangDiT .
Figure 3: Transformer architecture and self-conditioning in LangDiT . (a) The denoiser is read from bottom to top, with one of 12 blocks expanded. RMSNorm precedes each modulated sublayer; headwise query–key RMSNorm precedes RoPE, while values undergo neither operation. The 256-dimensional noise embedding supplies modulation parameters and residual gates (dashed arrows). A gated likelihood shortcut bypasses the Transformer and its untied vocabulary head. Only the shortcut’s explicit codebook factor is detached; C=sE remains trainable through noisy-state construction. (b) A bootstrap call to the same denoiser produces token probabilities without gradient tracking. These probabilities determine a clean-state estimate and an entropy-based weight for training self-conditioning (Equation ( 13 )). Inference passes the unweighted estimate to the next step. Shortcut warmup is omitted for clarity.
Component
Plaid baseline
Plaid + LangDiT
Backbone / head norm
LayerNorm
RMSNorm
Query–key norm
None
Headwise RMSNorm
Condition width
128
256
Shortcut gradients
Codebook factor detached on self-conditioned training examples
Same native rule
Shortcut calibration
Analytic coefficient
Same; no learned multiplier
Training self-conditioning
Unweighted bootstrap mean
Entropy-weighted bootstrap mean
Table 2: Scope of the Plaid adaptation. Four of the six LangFlow interventions are transferred; the native likelihood interface and gradient rules are retained.
OWT ( L=1024 )
Model
Train Steps
NFE
GenPPL ↓
Entropy
Published results: Autoregressive
Transformer
–
–
35.9
–
Published results: Absorbing-state diffusion
D3PM Absorb
–
–
–
–
DiffusionBert
–
–
–
–
Table 3: Generation results on OWT ( L=1024 ). Published baseline results are taken from Table 2 of Chen et al. (2026) and retain their original protocols; local results use our implementations, 1,024 samples, and EMA weights. Checkpoint denotes completed optimizer updates, and NFE counts denoiser evaluations. Extended-training checkpoints come from 1M-update training plans. Entropy denotes mean per-sample unigram entropy (Section 4.1 ). A dash indicates an unreported value; † denotes our implementation. For RePlaID, K+1 counts K reverse transitions and the final denoiser call.
Cumulative architecture
50k steps 128 NFE
170k steps 1024 NFE
Run
RMS
QK
c=256
Detach
Calib.
Train Conf.-SC
Transformer Params
GenPPL ↓
GenPPL ↓
Ent.
R1(baseline)
∼ 100M
87.48
41.97
5.24
R2
✓
∼ 100M
84.41
-
-
R3
✓
✓
∼ 100M
84.37
-
-
R4
✓
✓
✓
∼ 100M
80.74
-
-
R5
✓
✓
✓
✓
∼ 100M
79.35
-
-
Table 4: Cumulative architectural ablation on OWT-1024, from R1 to R7 (the complete model). Each successive row adds one component, with wider conditioning before shortcut detachment. The 50k-step column reports unsmoothed GenPPL using 128 NFE, 1,024 generated samples, and temperature 1. The 170k-step column reports unsmoothed GenPPL using 1024 NFE, 1,024 generated samples, and temperature 1.
Appendix figures & tables2 assets
Supplementary material from the paper’s appendix.
Appendix
Item
Value
Data
OpenWebText; document BOS/EOS markers; packed random crops
Table 5: Shared configuration and evaluation protocols for the cumulative architectural ablation. Sample counts and actual checkpoint steps accompany the reported measurements.
Configuration
Change relative to the preceding row
R1
Baseline
R2
RMSNorm
R3
Query–key RMSNorm
R4
256-dimensional noise conditioning
R5
Likelihood-shortcut detachment
R6
Learned likelihood-shortcut calibration
Appendix
Table 6: Cumulative architectural changes. Each row retains the changes introduced in preceding rows.
LUMIA Lab, School of Artificial Intelligence, Shanghai Jiao Tong University · School of Computer Science, Shanghai Jiao Tong University · Shanghai Innovation Institute +1