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, the authors 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.

Why Off-the-Shelf Transformers Fail Diffusion Text Generation

CDLMs generate text through iterative denoising rather than next-token prediction. This shift changes the computational profile in three ways that standard transformer designs do not address. First, the noise schedule introduces a time-varying signal-to-noise ratio that static normalization schemes cannot track. Second, gradient must flow through multiple denoising steps during training, creating longer effective computation graphs than autoregressive models. Third, the model must condition on noise level at every layer, yet most architectures borrow noise-conditioning patterns from image DiTs where spatial structure differs from sequential text. The result is a mismatch: transformers built for fixed-distribution inputs or 2D images produce unstable attention maps and degraded perplexity when pressed into CDLM service. The authors show this empirically by comparing vanilla LLM transformers and image DiTs against their proposed design on identical diffusion objectives.

Language Diffusion Transformer: The LangDiT Architecture

LangDiT is not a single trick but a suite of four coordinated modifications. Each targets a specific failure mode observed when standard transformers run CDLM training. The modifications are: feature and query-key normalization to stabilize activations across noise levels; wider noise conditioning to give every layer explicit access to the current diffusion timestep; likelihood-shortcut gradient routing to preserve signal through deep denoising chains; and confidence-gated self-conditioning to prevent early noisy predictions from corrupting later steps. Together they form a backbone that drops GenPPL by double-digit percentages across NFE budgets without changing the diffusion formulation, loss function, or sampler.

Feature and Query-Key Normalization in LangDiT

Standard LayerNorm normalizes across feature dimension assuming stationary statistics. In CDLMs, the distribution of hidden states shifts as noise anneals from pure Gaussian toward clean text embeddings. LangDiT replaces LayerNorm with a feature normalization that maintains running estimates of mean and variance per diffusion timestep, updated during training. This adapts the normalization to the current noise level. Separately, query-key normalization applies a cosine-similarity attention mechanism with learned temperature per head. By normalizing queries and keys before the dot product, attention weights become invariant to the overall scale of hidden states, which grows as denoising progresses. The combination prevents attention collapse in late denoising steps where clean-signal magnitude dominates noisy residuals.

Wider Noise Conditioning Mechanisms

Most CDLMs inject noise level through a single embedding added to the input or a FiLM-style modulation at one layer. LangDiT broadcasts a sinusoidal timestep embedding to every transformer block and applies it through two parallel pathways: an additive bias to the attention output and a multiplicative gate on the feed-forward output. This wider conditioning ensures each layer can adjust its computation to the current denoising stage independently. The authors ablate this design against single-point injection and find consistent perplexity improvements, especially at low NFE counts where each step carries more denoising burden.

Likelihood-Shortcut Gradient Routing and Calibration

During training, the diffusion loss computes gradients at the final denoising step, which must propagate back through all intermediate steps. In deep denoising chains, this gradient attenuates. LangDiT adds an auxiliary likelihood head at each layer that predicts the clean token distribution from that layer's hidden state. A weighted sum of these per-layer likelihood losses provides a shortcut gradient path to early layers. The weights are calibrated by a learned schedule that emphasizes early-layer shortcuts at high noise and late-layer shortcuts at low noise. This calibration matches the information geometry of the diffusion process: early layers see mostly noise and benefit from strong likelihood signals, while late layers refine clean structure and need finer gradient control.

Confidence-Gated Training Self-Conditioning

Self-conditioning feeds the model's own previous prediction as an additional input to the next denoising step. In image diffusion this stabilizes sampling, but in text the discrete token space makes early predictions unreliable. LangDiT gates self-conditioning with a confidence score derived from the entropy of the model's predicted token distribution at the previous step. Low entropy (high confidence) enables the self-conditioning pathway; high entropy suppresses it. This prevents error propagation when the model is uncertain, while still exploiting self-consistency when predictions are reliable. The gate is differentiable and trained end-to-end, requiring no threshold tuning.

Experimental Results: LangFlow and Plaid Formulations

The authors evaluate on two CDLM families. LangFlow uses a flow-matching formulation with continuous noise schedules. Plaid uses a discrete diffusion formulation with absorbing states. On OpenWebText 1024-token unconditional generation, LangDiT-LangFlow achieves GenPPL 21.72 at 1024 NFEs and 44.72 at 128 NFEs. LangDiT-Plaid achieves 11.33 at 1024 NFEs and 27.51 at 128 NFEs. Both represent new best published numbers for their respective formulations. The gap between 128 and 1024 NFE scores narrows under LangDiT, indicating more efficient denoising per step. Ablations confirm each of the four components contributes; removing any single one degrades perplexity by 2 to 5 points depending on NFE budget.

Why Architectural Design Matters More Than Objective Changes

A central claim of the paper is that transformer architecture is a first-order factor in CDLM performance. To isolate this, the authors keep representation (continuous embeddings), training objective (score matching or flow matching), and inference sampler (ODE solver or discrete transition) fixed while swapping only the transformer backbone. The perplexity drops attributable to LangDiT match or exceed gains reported from changing diffusion formulations. This suggests the field has underinvested in architecture search for diffusion text. The authors argue that future CDLM progress should treat architecture co-design with the diffusion process as standard practice, not afterthought.

Practical Implications for Developers Building CDLMs

LangDiT is designed as a drop-in replacement for the transformer backbone in existing CDLM codebases. The normalization layers, noise conditioning pathways, likelihood heads, and confidence gate add modest parameter overhead (approximately 5 percent over a vanilla transformer of equivalent width and depth) and integrate with standard PyTorch or JAX training loops. Developers can adopt individual components incrementally: feature normalization and wider noise conditioning require the least refactoring; likelihood shortcuts and confidence gating need loss-function modifications but no sampler changes. The paper releases reference implementations for both LangFlow and Plaid integration points. For teams training CDLMs at scale, the NFE efficiency gains translate directly to inference speedups or quality improvements at fixed compute budgets.

Limitations and Open Questions

The study focuses on unconditional generation on OpenWebText. Conditional tasks (prompted generation, infilling, translation) and longer context lengths remain unevaluated. The confidence gate relies on entropy of predicted distributions, which may behave differently under temperature scaling or top-k sampling at inference. The likelihood shortcut introduces additional memory cost during training proportional to layer count times vocabulary size; the authors use a sampled softmax approximation but do not report its precision trade-off at scale. Finally, the architectural search space is constrained to modifications of the standard transformer block; alternative architectures (state-space models, linear attention, hybrid designs) are not explored. The authors note that LangDiT principles may transfer to those families but leave experimental validation to future work.

Read the paper on arXiv