PaperlayerView exact revision on arXiv

Source-linked AI summary

Designing Transformers for Continuous Diffusion Language Models

Weijian Luo, Yujian Chen, Hao Jiang, Guang Cheng, Tianyang Hu

arXiv:2610.22330v1cs.AI

TL;DR

Continuous diffusion language models offer parallel generation, but their Transformer architectures have received less systematic attention and commonly inherit designs from causal LLMs or image DiTs. This paper studies architecture within fixed CDLM formulations and introduces LangDiT. In matched LangFlow comparisons, LangDiT lowers GenPPL, while the authors identify architecture as a reusable source of improvement across formulations.

  • Problem

    Transformer architectures for CDLMs have received less systematic attention, and existing models commonly inherit designs from causal LLMs or image-generation DiTs despite CDLMs’ distinct denoising requirements.

    Source: §1 · pp. 1, 3

  • Method

    The authors systematically study CDLM Transformer design and introduce LangDiT, varying architecture within fixed formulations while retaining each model’s representation, objective, and inference sampler.

    Source: §1 · p. 3

  • Results

    At 128 NFE, matched LangFlow GenPPL falls from 64.26 to 53.66, a relative reduction of 16.5%.

    Source: §4.2 · p. 13

  • Takeaways & Limitations

    The findings identify architectural design as an important, reusable source of improvement in continuous diffusion language models.

    Source: §1 · p. 3, §5 · p. 16

    The study is limited to OWT-1024 at one backbone scale and one training seed, and GenPPL and unigram entropy capture limited aspects of generation quality.

    Source: §4.4 · p. 16

Abstract

from arXiv · show

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.

1. Introduction

Continuous diffusion language models refine continuous representations in parallel, but their Transformer backbones have received less systematic attention than their formulations and learning processes. This paper studies architecture within fixed CDLM formulations and introduces LangDiT, reporting lower GenPPL in matched LangFlow comparisons.

1. Introduction

  • The paper targets a gap in systematic study of Transformer architecture for CDLMs, whose models commonly inherit designs from causal LLMs or image-generation DiTs.CDLMs recover discrete tokens from noisy continuous states while coordinating noise-dependent computation, vocabulary prediction, and sometimes self-conditioning.Source: §1 · pp. 1, 3
  • The study evaluates unconditional OpenWebText generation at sequence length 1,024 and retains each formulation’s representation, objective, and inference sampler when changing architecture.LangDiT is developed within LangFlow and transferred to a Plaid-style CDLM, reusing compatible components while retaining its native prediction interface.Source: §1 · p. 3
  • At 128 NFE, the seven-stage LangFlow ladder reduces GenPPL from 64.26 to 53.66 at 250k updates; extended-training scores do not isolate architecture from training budget.The extended-training LangFlow and Plaid checkpoints reach GenPPL 21.72 and 11.33 at the larger sampling budget.Source: §1 · p. 3
  • LangDiT combines backbone normalization and noise conditioning with likelihood-shortcut gradient routing and calibration, plus confidence-gated training self-conditioning.The study examines feature normalization and noise conditioning, likelihood-shortcut gradients, and the strength of self-conditioning under shared training and evaluation settings.Source: §1 · p. 3
  • The authors frame architectural design as an important source of progress in CDLMs and a reusable foundation for further research on parallel language generation.Their contribution includes a cumulative LangFlow ablation and examination of compatible components in a Plaid-style formulation.Source: §1 · p. 3

2. Related Work

Related work spans continuous denoising and flow, latent representations, Transformer normalization and conditioning, self-conditioning, and discrete diffusion evaluation. This study focuses on architectural design within fixed LangFlow and Plaid-style formulations.

Continuous language diffusion and flow

  • Prior continuous-language work spans denoising, categorical diffusion, likelihood-based training, flow maps, learned embeddings, hyperspherical dynamics, and embedding-space flow matching.The paper studies architectural design within fixed LangFlow and Plaid-style formulations.Source: §2 · p. 4

Latent representations and decoding

  • Latent-language models vary their decoding and representation designs, while this work changes components within the existing codebook-based prediction interface.Examples include contextual autoregressive decoding, jointly learned encoder-denoiser-decoder systems, text VAEs, and restricted noisy-input pathways.Source: §2 · p. 4

Normalization and conditioning

  • Prior diffusion and Transformer work studies adaptive conditioning and magnitude control, and this paper examines normalization choices alongside noise-conditioning capacity for language denoising.The cited approaches include RMSNorm, query-key normalization, and work on attention-entropy collapse.Source: §2 · p. 4

Self-conditioning

  • Earlier systems use previous predictions to guide denoising; here, an entropy-based gate regulates bootstrap inputs during training while retaining the baseline loss and inference sampler.The related work also includes trajectory consistency as an approach to training-inference mismatch.Source: §2 · p. 4

Discrete diffusion and evaluation

  • Because likelihood rankings can differ from generation quality and sampling cost, the paper reports GenPPL alongside mean per-sample unigram entropy and explicit sampling budgets.Prior work also notes that lower GenPPL can accompany reduced distributional coverage.Source: §2 · p. 4

3. The Language Diffusion Transformer

LangDiT is a Transformer denoiser for continuous language generation, developed within a LangFlow-style baseline while retaining its representation, objective, and sampler. Its design changes the denoising backbone, vocabulary shortcut, and training-time self-conditioning, with compatible components also adapted to a Plaid-style formulation.

3.1. Overall Denoising Pipeline

  • At each denoising step, a bidirectional Transformer processes noisy embeddings and optional self-conditioning, then combines contextual logits with a likelihood shortcut to estimate the clean state.The probability-weighted embedding mean supplies the continuous estimate for sampler updates and subsequent self-conditioning.Source: §3 · p. 4, §3.1 · p. 5

3.3. Noise-Conditioned Transformer Backbone

  • Backbone RMSNorm, headwise query–key RMSNorm, and wider noise conditioning target feature scale, attention-logit scale, and noise-dependent modulation capacity.The conditioning width increases from 128 to 256 without changing Transformer hidden width or feed-forward dimensions; it adds 7.36M parameters, or 2.25% of baseline.Source: §3.1 · p. 5, §3.3 · p. 8

3.4. Vocabulary Prediction with a Likelihood Shortcut

  • The likelihood shortcut combines contextual logits with Gaussian-corruption evidence, while stop-gradient blocks its direct codebook-gradient path and a bounded learned gate calibrates its forward contribution.The codebook remains trainable through noisy-state construction, so detachment changes the optimization graph without changing the forward model class or adding parameters.Source: §3.4 · p. 9

3.5. Self-Conditioning during Training and Sampling

  • During training, an entropy-based confidence weight attenuates the bootstrap clean-state estimate; inference retains the unweighted posterior mean for self-conditioning and sampler updates.The weight is zero for uniform bootstrap probabilities and one for a one-hot distribution; 25% of training examples use bootstrap self-conditioning.Source: §3.5 · pp. 9–10

3.6. Adaptation to a Plaid-Style Formulation

  • The Plaid-style adaptation transfers four of six LangFlow interventions while retaining the native likelihood interface and gradient rules.It preserves the 16-dimensional state, NELBO, learned schedule, and DDPM transitions; the experiment tests reuse of compatible components, not transfer of the entire LangFlow pipeline.Source: §3.6 · pp. 10–11, Table 2 · p. 11

3.7. Implementation Details

  • OpenWebText LangDiT retains baseline Transformer depth and hidden width, using 12 blocks, width 768, 12 heads, and feed-forward width 3072.The cumulative architecture ladder adds backbone RMSNorm, Q/K RMSNorm, wider conditioning, shortcut detachment, learned calibration, and training-time confidence gating in separate runs from scratch.Source: §3.7 · p. 11, §3.8 · p. 11

3.8. Architecture Variants

  • The cumulative ablation defines LangDiT in its final row, with about 100M Transformer parameters including the codebook and untied vocabulary projection.The ladder uses a 250k-update schedule, while extended-training runs are reported separately; baseline and LangDiT also share fused QKV projections, rotary position embeddings, GELU feed-forward layers, and zero-initialized adaptive-normalization projections.Source: §3.8 · p. 11, §3.7 · p. 11

Text Generation Sample of RePlaID + LangDiT

4. Experiments

Experiments compare LangDiT with prior models and use shared-schedule LangFlow ablations to assess architectural changes, then report extended-training results, Plaid-style transfer, and evaluation boundaries.

4.1. Experimental Setup

  • All seven LangFlow configurations shared data order, objective, schedule, and sampler; evaluation used unconditional 1,024-token samples and GPT-2 Large scoring.GenPPL excludes each sample's first token, while unigram entropy is a lexical statistic that does not establish semantic diversity or distributional coverage.Source: §4.1 · pp. 12–13

4.2. Generation Results

  • Published OWT results report GenPPL 36.5 for LangFlow and 35.9 for the autoregressive Transformer, retaining their original protocols.The table provides a broad comparison; the paper attributes architectural differences using its shared-schedule local comparison.Source: §4.2 · p. 13
  • Under the shared 250k-update schedule, LangDiT reduced GenPPL from 64.26 to 53.66 at 128 NFE, a 16.5% relative reduction.Mean per-sample unigram entropy changed from 5.416 to 5.391; data order, batch size, objective, and sampler were held fixed.Source: §4.2 · p. 13
  • GenPPL reached 21.72 at 1,024 calls for LangFlow and 11.33 at 1,025 total calls for Plaid; extended checkpoints do not isolate architecture.At lower budgets, the corresponding scores were 44.72 at 128 calls and 27.51 at 129 total calls; both checkpoints came from 1M-update plans but were not completed 1M-step evaluations.Source: §4.2 · p. 13, Table 3 · p. 14

4.3. Cumulative Ablation and Analysis

  • Backbone RMSNorm reduced 50k-step GenPPL from 87.48 to 84.41; widening noise conditioning reduced it further to 80.74, but added 7.36M parameters.The Q/K RMSNorm comparison was small at 50k steps but larger at 250k; the conditioning result cannot isolate the benefit of capacity allocation.Source: §4.3 · p. 15
  • At 50k updates, R1–R7 reduced GenPPL by 29.4%, with confidence-gated training self-conditioning producing the largest adjacent reduction, from 74.03 to 61.73.These increments depend on earlier modifications, are not independent effects, and do not establish statistical significance from a single seed.Source: §4.3 · p. 15, Table 4 · p. 15

4.4. Parameters and Computational Cost

  • The Plaid-style adaptation reuses compatible LangDiT components while retaining a different continuous state space, objective, and sampler.The reported Plaid scores do not provide a matched baseline comparison for all six LangFlow interventions.Source: §4.2 · p. 13, §4.4 · p. 16
  • The study is limited to OWT-1024, one backbone scale, and one seed; equal updates and denoiser calls do not establish equal FLOPs or wall-clock time.The paper calls for parameter- and compute-matched comparisons, component-removal ablations, more seeds, and broader quality evaluation.Source: §4.4 · p. 16

5. Conclusion

LangDiT separates feature-scale control, noise modulation, shortcut gradients, and self-conditioning inputs. On OWT-1024, it improves generation performance within LangFlow and transfers compatible components to a Plaid-style model while preserving each formulation's representation, objective, and sampler.

5. Conclusion

  • LangDiT improves generation performance in LangFlow and adapts to a Plaid-style model while preserving each model's representation, objective, and inference sampler.The study evaluates GenPPL alongside unigram entropy to interpret performance gains together with lexical-diversity changes.Source: §5 · p. 16
  • LangDiT separates feature-scale control from noise modulation, direct codebook-gradient routing from the shortcut's forward contribution, and clean-state estimation from self-conditioning-input weighting.Source: §5 · p. 16
  • The findings identify architectural design as a reusable source of improvement in CDLMs and motivate evaluation at larger scales and across broader generation tasks.Source: §5 · p. 16

A. Architectural Analysis and Derivations

The appendix derives identities used to motivate LangDiT's architectural interventions. Its analysis does not establish convergence or generation-quality guarantees.

A. Architectural Analysis and Derivations

A.1. Scale Dependence of Normalized Attention

The appendix analyzes normalized attention, likelihood evidence, shortcut-gradient detachment, and posterior-mean self-conditioning. It states limits on what these identities imply about optimization and uncertainty.

A.1. Scale Dependence of Normalized Attention

  • For fixed query and key gains, normalized attention bounds the dot-product contribution to logits before softmax, but additive attention bias must be considered separately.With ε > 0, positive-rescaling invariance is approximate when ε is negligible relative to the inputs' mean squared magnitudes; learned gains can still grow and attention may remain concentrated.Source: §A.1 · p. 20

A.2. Observation Evidence and Contextual Prediction

  • Bayes' rule separates token odds into local observation evidence and evidence from other noisy positions plus the token prior, without requiring a uniform prior.The learned vocabulary head receives the noisy state at the target position too; forward equivalence between gated and ungated shortcut representations does not imply equivalent gradients or optimization dynamics.Source: §A.2 · p. 21

A.3. The Gradient Path Removed by Detachment

  • Detaching the shortcut's explicit codebook factor removes its direct gradient contribution while leaving other paths in the computational graph intact.This identity does not determine whether the removed contribution helps or hinders optimization; a large shortcut coefficient alone does not imply a large gradient.Source: §A.3 · pp. 21–22

A.4. Posterior Means, Geometry, and Self-Conditioning

The equal-norm codebook further implies

  • Entropy does not determine embedding variance, so entropy-based self-conditioning weights are a heuristic for auxiliary-input weighting rather than an exact embedding-space uncertainty measure.In LangDiT, confidence weighting changes only the training bootstrap input; inference uses the unweighted posterior mean.Source: The equal-norm codebook further implies · p. 22

B. Implementation Configuration

The implementation section specifies the shared ablation configuration and evaluation protocols, along with selected architectural settings for the cumulative variants.

B. Implementation Configuration

  • The 250k-update budget is a configured training target, while each measurement is tied to its recorded checkpoint step.The shared-schedule ablation and separate extended-training checkpoints are reported separately.Source: B. Implementation Configuration · p. 22
  • Evaluation uses 1,024 samples of length 1,024 at 128 NFE during training, with additional checkpoint evaluation at 128 or 1,024 denoising steps.GenPPL uses a frozen GPT-2 Large evaluator with corpus-level token aggregation.Source: B. Implementation Configuration · p. 23
  • For variants R3–R7, query–key RMSNorm follows projection and head reshaping and precedes RoPE, with separate 64-dimensional query and key gains shared across heads within each layer.Source: B. Implementation Configuration · p. 23
  • Entropy-based confidence weighting appears only in R7 and affects the bootstrap input for self-conditioned training examples, while inference retains the unweighted posterior mean.Source: B. Implementation Configuration · p. 23

B.1. Separate Extended-Training Runs

The extended LangFlow and Plaid-style evaluations use separate training plans and continuation schedules from the 250k ablation ladder. Their selected checkpoint steps do not mean the 1M-update targets were reached.

B.1. Separate Extended-Training Runs

  • The extended LangFlow + LangDiT evaluation uses a 660k checkpoint after continued training from step 450k under a separate 1M-update plan.Its cosine decay is indexed from step 270k; the 250k ablation schedule is separate.Source: §B.1 · p. 23
  • The extended Plaid + LangDiT evaluation uses a 520k checkpoint, with continuation from step 250k under a 1M-update target plan.The continuation uses warmup followed by cosine decay; both training stages use batch size 512 and EMA decay 0.999.Source: §B.1 · p. 23
  • All extended-training results in Table 3 use plain, unweighted inference self-conditioning; release metadata separately records source run, actual checkpoint step, planned target, and sampling mode, and selected results do not imply 1M-update targets were reached.Source: §B.1 · p. 23

C. Ablation Configurations

The ablation configurations form a cumulative ladder, with each configuration retaining preceding changes. R7 is the full LangDiT architecture under a 250k-update budget, distinct from the separate 1M-update experiment.

C. Ablation Configurations

  • R7 is the full LangDiT configuration with a 250k-update training budget; the 1M-update configuration is a separate experiment.Table 6 summarizes the seven cumulative configurations used in Table 4 and the Figure 1 training trajectories.Source: C. Ablation Configurations · p. 24
  • R4 widens noise conditioning to 256 dimensions, and R5 adds shortcut detachment at that width.Each row of Table 6 retains changes introduced in preceding rows.Source: C. Ablation Configurations · p. 24
  • The Table 4 50k-step GenPPL values are readings at optimizer step 50,000, not minima or averages over earlier checkpoints.The 250k Table 3 comparison uses R1 and R7 under the same 250k schedule.Source: C. Ablation Configurations · p. 24

D. Evaluation and Reporting Details

The evaluation reports matched generation settings, defines its generation and entropy metrics, and documents sample accounting and scorer provenance. It distinguishes the training-time confidence-weighted bootstrap input from the baseline inference self-conditioning rule.

D. Evaluation and Reporting Details

  • Matched ablations use the same checkpoint steps, EMA conventions, generation seeds, temperatures, NFE, and sample counts, with 1,024-token generations.Logged trajectories and additional checkpoint evaluations use 1,024 samples.Source: D. Evaluation and Reporting Details · p. 24
  • GenPPL is the exponential of mean negative log-likelihood over scored tokens, while unigram entropy is averaged across samples using each sample’s own token-frequency distribution.The sample distributions are not pooled into one frequency distribution.Source: D. Evaluation and Reporting Details · p. 24
  • Checkpoint metrics come from files with 1,024 generated records; texts shorter than two evaluator tokens contribute no next-token predictions to GenPPL.Entropy is computed on non-padding tokens; the empty-clean-text fallback is described in Section 4.1.Source: D. Evaluation and Reporting Details · p. 24
  • Training-curve and 50k-step results use logged R1–R7 scalars with EMA weights, 1,024 samples of length 1,024, 128 NFE, temperature 1, and generation seed 0.The generation seed differs from the training seed of 42, and entropy uses natural logarithms.Source: D. Evaluation and Reporting Details · pp. 24–25
  • The exporter preserves optimizer-step indices and raw scalar precision, rounds only displayed values, and does not impute missing TensorBoard measurements.Reported logs retain actual checkpoint steps and evaluation settings.Source: D. Evaluation and Reporting Details · p. 25
  • R7 weights the training bootstrap input but retains baseline unweighted self-conditioning at sampling; inference-time weighting defines a different sampling setting.Model class, code revision, sampler implementation, and feedback rule are needed to identify an evaluation unambiguously.Source: D. Evaluation and Reporting Details · p. 25

E. Limitations and Further Controls

The reported limitations concern what the cumulative ablations can isolate, missing controls for the two gates, and the broader settings needed to assess robustness. Potential failure modes also qualify conclusions drawn from the current configurations.

E. Limitations and Further Controls

  • The cumulative ladder does not establish each component’s contribution in the full model or resolve component interactions.Suggested controls include removing components from the full model and crossing selected normalization, conditioning-width, detachment, and calibration settings.Source: E. Limitations and Further Controls · p. 25
  • The reported results omit controls that would distinguish noise-dependent shortcut calibration from simpler multipliers and token-specific confidence gating from generic attenuation.Proposed comparisons include fixed or learned constants and confidence gates shuffled within noise bins.Source: E. Limitations and Further Controls · p. 25
  • Assessing training variability would require independent baseline and modified-model runs across shared seeds; resampling generated texts does not capture that uncertainty.Broader comparisons also require matched data exposure and optimization schedules, comparable tuning budgets, and parameter and FLOP accounting.Source: E. Limitations and Further Controls · p. 25
  • Potential failure modes include suppressing useful uncertain self-conditioning inputs or removing shortcut gradients that aid optimization.These possibilities limit how broadly conclusions from the present configurations generalize.Source: E. Limitations and Further Controls · p. 25

F. Qualitative Generations

The section presents fixed-ID generations from the evaluations, with settings tied to Figure 2. The examples illustrate those evaluations but are not evidence of typical quality or systematic improvement with NFE.

F. Qualitative Generations

  • Figure 2’s examples cover baseline and modified evaluations, plus extended-training evaluations; Table 3 reports corresponding checkpoints and results.RePlaID pairs in panel (a) differ in checkpoint progress, and its modified 250k checkpoint belongs to a 1M-update plan.Source: F. Qualitative Generations · p. 26
  • Fixed sample IDs are shown from twelve evaluation files, and each card’s GenPPL describes its full 1,024-sample evaluation rather than the displayed text alone.Paired entropy is reported in Table 3.Source: F. Qualitative Generations · p. 26
  • LangFlow examples use Euler–EDM with 128 or 1,024 denoiser calls; RePlaID uses 128 or 1,024 reverse transitions plus one final call.Examples use EMA weights, 1,024 model tokens, temperature 1, final argmax decoding, and unweighted inference self-conditioning.Source: F. Qualitative Generations · p. 26
  • The fixed-ID examples illustrate the cited evaluations but do not establish typical quality, factual accuracy, or systematic improvement with NFE.Text is reproduced without correction or truncation, including special tokens and document boundaries.Source: F. Qualitative Generations · p. 26

F.1 LangFlow baseline

The supplied baseline examples include raw model outputs labeled by checkpoint, EMA status, seed, temperature, and sequence length. Their text includes a court-case account and passages on border policy and economic issues.

F.1 LangFlow baseline

  • The example labels specify unedited output at a 250k checkpoint, with EMA weights, temperature 1, 1,024 model tokens, and seeds 1 and 2.The labels appear in separate sample headers.Source: §F.1 · pp. 27–28
  • Other output passages discuss immigration and border policy, workers’ wages, economic conditions, and tariffs on Chinese goods.The text includes discussion of border safety and U.S. workers, alongside claims about wages and trade.Source: §F.1 · p. 27
  • One output presents a court-case account involving Bobby Wise and Judge Featherstone, with later passages discussing trials, sentencing, and release.The supplied passages include accounts of proceedings and racial bias.Source: §F.1 · p. 28

Related papers

Loading 2610.22330v1…