PaperlayerView exact revision on arXiv

Source-linked AI summary

Less Uniform Discrete Diffusion is More Powerful and Scalable

Kaibo Wang, Ding Ding, Fangyu Ding, Zijin Feng, Han Shi, Haili Bai, Jiacheng Sun, Yang Xiang

arXiv:2609.35817v1cs.CLcs.AI

TL;DR

Uniform diffusion language models have shown promise, but scaling them to complex reasoning remains challenging. LUDI combines a less uniform loss with token-level corruption hints, and experiments report competitive 7B performance with faster generation than autoregressive decoding.

  • Problem

    UDLMs have shown advantages for few-step generation, but their viability at larger scales and on complex reasoning tasks remains unestablished.

    Source: §1 · p. 1

  • Method

    LUDI removes the loss term encouraging uniform predictions and gives each token a time embedding that indicates its probability of random replacement.

    Source: §1 · p. 2

  • Results

    LUDI-7B achieves approximately 3 tokens per step with comparable performance to its AR model and competitive performance with masked diffusion models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    LUDI-7B demonstrates that a UDLM continued from an AR checkpoint can handle complex reasoning, while the potential of UDLMs for complex generation remains largely untapped.

    Source: §1 · p. 2, §6 · p. 9

    The work has not fully activated UDLM self-correction, and computational constraints prevented pretraining a large-scale UDLM from scratch.

    Source: Limitations · p. 10

Abstract

from arXiv · show

Although uniform diffusion language models (UDLMs) represent a promising diffusion paradigm, scaling them remains challenging. We identify the core obstacle as an over-uniform training objective and condition-target confusion during sampling. To address these, we propose Less Uniform Diffusion (LUDI), a novel UDLM framework. Specifically, we (i) introduce a less uniform loss that directs each reverse transition toward the clean token, and (ii) equip the model with per-token time embeddings that supply token-level corruption hints, enabling confidence-based few-step sampling. Experiments across scales show that LUDI yields cleaner supervision and improves few-step generation. We further continue-train a 7B autoregressive model into LUDI-7B, resulting in a UDLM capable of complex reasoning. It achieves a 3-token-per-step speedup over AR decoding and competitive performance compared with masked diffusion baselines, revealing that the full potential of UDLMs for complex generation remains to be unlocked.

1 Introduction

Uniform diffusion language models have shown promise for few-step generation, but their ability to scale to complex reasoning remains unestablished. LUDI addresses over-uniform training and condition-target confusion, and experiments report competitive quality with faster generation at 7B scale.

1 Introduction

  • UDLMs use a uniform vocabulary prior and show potential for few-step generation, but their viability for complex reasoning and larger scales remains unestablished.Prior demonstrations were limited to small scales, small datasets, or scaling-law investigations.Source: §1 · pp. 1–2
  • UDLM scaling is hindered by over-uniform training signals and uniform denoising, which can confuse conditions with targets and cause repeated-token collapse on complex tasks.The training objective weakens clean-token confidence, while sampling fails to decode simpler conditioning tokens first.Source: §1 · p. 2
  • LUDI removes the loss term that encourages uniform predictions and adds per-token time embeddings indicating whether tokens were randomly replaced.The LU loss steers reverse transitions toward clean tokens; token-level hints help distinguish conditions from targets and enable confidence-based, easy-to-hard sampling.Source: §1 · p. 2
  • LUDI-7B, continued from an autoregressive checkpoint, performs complex reasoning and achieves about 3 tokens per step with comparable performance to its AR model and competitive performance with masked diffusion models.Experiments at 170M and 1B scales also find cleaner supervision and improved few-step generation and downstream task performance from the LU loss.Source: §1 · p. 2

2 Background

Uniform-state diffusion gradually corrupts clean tokens toward a uniform vocabulary prior, then learns a reverse process to denoise from that prior. The background describes the reverse-transition objective and related parameterizations and simplified losses.

2.1 Forward Process

  • Uniform-state diffusion gradually perturbs each clean token toward the uniform distribution over the vocabulary as time advances.The forward process is a continuous-time Markov process; the reverse process generates samples by denoising from the uniform prior.Source: §2.1 · p. 3, §2 · p. 3

2.2 Reverse Process and Training Objective

  • The exact reverse posterior provides the target for learning a parameterized reverse transition from each noisy token state.The standard objective minimizes a local KL divergence, which contributes to the negative evidence lower bound.Source: §2.2 · p. 3, §2.1 · p. 3
  • Duo parameterizes the reverse transition by predicting the clean token, while SDDLM introduces empirically simplified losses that are not theoretically NELBO objectives.The uniform-state SEDD objective with Duo’s x0 parameterization yields the same generator-level objective; SDDLM reports stronger empirical performance for its simplified losses.Source: §2.2 · pp. 3–4

3 Methodology

LUDI addresses the over-uniform supervision and condition-target confusion that hinder UDLM scaling with a less uniform loss and per-token time embeddings. Together, these changes support stronger few-step generation and scaling to complex tasks.

3 Methodology

  • LUDI combines a less uniform loss with per-token time embeddings to address over-uniform supervision and condition-target confusion in UDLMs.The method aims to decode high-confidence tokens quickly and use them as conditioning context for complex generation.Source: §3 · p. 4

3.1 Less Uniform Loss

  • The ELBO’s mode-covering KL objective encourages overly uniform predictions, leaving little effective supervision toward the clean token at large t.Under the stated common schedule, only half of the probability mass is placed on the clean token on average.Source: §3.1 · p. 4
  • LU loss removes the term that encourages uniform predictions, avoiding excessive smoothing while aligning the parameterized transition rate with the clean token.The paper says this stronger signal improves few-step generation through consistency.Source: §3.1 · p. 5
  • At corrupted positions, LU loss reduces the probability of the incorrect token and increases the probability of the clean token, encouraging transitions toward the clean token at any time step.The paper also says its formulation avoids singular logarithmic terms.Source: §3.1 · p. 5
  • With optimized operators, LUDI achieves a 4.39× speedup and a 3.45× memory reduction over the ELBO loss.The loss involves only indices y and m, whereas the ELBO also transforms and reduces over all vocabulary entries.Source: §3.1 · p. 5

3.2 Scaling up UDLM

  • UDLMs provide only global time, so the model struggles to distinguish condition tokens from denoising targets and tends to denoise all positions simultaneously.The paper identifies this ambiguity as particularly harmful for complex reasoning tasks.Source: §3.2 · p. 5
  • Per-token time embeddings assign each position a distinct time value indicating its probability of random replacement, providing a hint about whether it is a condition or target.The forward process samples per-token times with mean equal to global time, preserving the marginal distribution of x_t.Source: §3.2 · p. 5
  • Conditional sampling progressively commits high-confidence tokens as conditions, and the paper reports that this easy-to-hard denoising schedule is critical for complex reasoning tasks.The method initializes from an autoregressive model to preserve acquired knowledge when continuing training as a block-UDLM.Source: §3.2 · p. 6

4 Experiments

Experiments span 170M and 1B pretraining and 7B AR-to-UDLM continue-training, evaluating generation, downstream performance, and speed. LUDI performs strongly across these settings, while practical speedups remain constrained by batch sampling and infrastructure overhead.

4 Experiments

  • The experiments cover 170M and 1B small-scale pretraining, plus 7B continue-training from an AR checkpoint for complex-reasoning evaluation.The small-scale models use OpenWebText at 170M and FineWeb at 1B; LUDI-7B is initialized from Qwen2.5-7B-Instruct and trained on Dolci-Instruct-SFT.Source: §4 · p. 6, §4.1 · p. 6, §4.2 · p. 7

4.1 Small-Scale LUDI

  • At 1024 sampling steps, LUDI reaches a generative perplexity of 37.76 and entropy of 7.50, outperforming UDLM and MDLM baselines.Across most budgets of at least 64 steps, LU loss yields stronger generation quality; the reported entropy reduction does not stem from mode collapse.Source: §4.1 · p. 6, Figure 3 · p. 7
  • LUDI attains the highest average score and leads on most downstream tasks among the evaluated 1B models.At 200K training steps, LUDI already outperforms all baselines across downstream tasks, which the authors attribute to the LU loss’s cleaner supervision and direct matching to x0.Source: §4.1 · p. 7, Table 2 · p. 7

4.2 LUDI-7B

  • LUDI-7B scores 65.3 on average, outperforming the baselines and its AR initialization across the reported evaluation.It performs strongly on MMLU and GPQA but underperforms on mathematical reasoning, which the authors attribute mainly to training-data distributional bias.Source: §4.2 · p. 8, Table 3 · p. 8
  • At thresholds above 0.7, LUDI-7B maintains stable GSM8K accuracy while achieving a 2.6–3.0 token-per-step speedup over AR decoding.The token-per-step figure is theoretical; measured speedups at small batch sizes are 1.30–1.92×, with larger-batch speedups requiring further infrastructure optimization.Source: §4.2 · p. 8, Figure 4 · p. 8
  • Without per-token time embeddings, LUDI-7B collapses to repeating a single token; replacing LU loss with SDv1 also causes a clear performance drop.The ablations also identify the AR loss and causal context attention as important, and LUDI-7B surpasses FDNB trained under the same budget.Source: Table 3 · p. 8, §4.2 · p. 9

5 Related Works

Related work contrasts masked diffusion models’ strong complex-reasoning results with UDLMs’ potential advantages and remaining scale limitations. Prior UDLM work has not yet established complex-reasoning performance at scale or directly addressed over-uniform training and position-uniform sampling together.

Masked diffusion language models

  • Masked diffusion models have reached performance comparable to frontier LLMs on complex reasoning while offering inference-speed advantages.Their sampling has moved toward confidence- and entropy-based policies, but the binary mask/non-mask state space makes correcting generated tokens difficult and the mask state constrains parallel decoding.Source: §5 · p. 9

Uniform diffusion language models

  • UDLMs use uniform-vocabulary corruption rather than an absorbing mask, with proposed advantages in few-step generation, controllability, self-correction, and data- or compute-limited scaling.Existing analyses remain largely small-scale or focused on scaling laws; prior efforts have not scaled UDLMs for complex reasoning and often leave over-uniform objectives or position-uniform sampling unaddressed.Source: §5 · p. 9

6 Conclusion

The conclusion presents LUDI as a step toward practical UDLM scaling, attributing the main barrier to excessive uniformity in training objectives and sampling rather than the uniform noise prior itself. LUDI combines a less uniform loss with token-level time awareness, and LUDI-7B demonstrates extension to complex reasoning tasks.

6 Conclusion

  • LUDI addresses excessive uniformity in UDLM training and sampling with a less uniform loss and token-level time awareness, and LUDI-7B extends the paradigm to complex reasoning.The authors characterize UDLMs as a promising but underexplored direction.Source: §6 · p. 9

Limitations

The paper identifies limits of LUDI’s current scope, while the appendix derives uniform-state diffusion preliminaries and theoretical results motivating the LU loss.

Limitations

  • LUDI has not yet fully activated self-correction, and the authors say teaching models to detect and amend their own mistakes may require dedicated post-training.They also have not pre-trained a large-scale UDLM from scratch; continuing from an AR checkpoint preserves knowledge but may retain AR inductive biases.Source: Limitations · p. 10

A Proofs and Derivations

A.1 Uniform-State Diffusion Preliminaries

A.1.3 SEDD–Duo Reparameterization

  • Under the x0-parameterization, the SEDD score-entropy objective reduces exactly to the Duo objective.Source: §A.1.3 · p. 14

A.1.4 Contrast with Mask Diffusion

  • For absorbing-state mask diffusion, the learned posterior KL reduces to clean-token cross-entropy because both posterior weights are analytic.Source: §A.1.4 · p. 14

A.2 ELBO Loss and LU Loss

  • The appendix establishes that the LU contrast is the finite model-dependent part of a small-step Dirac-target KL and approximates the uniform-state ELBO integrand under stated assumptions.The ELBO approximation concerns corrupted positions and has an o(1) remainder as vocabulary size grows; already-clean positions contribute uniformly o(1) under an additional assumption.Source: §A.2 · p. 14, §A.2.1 · p. 15

A.2.2 Proof of Theorem 1

The appendix relates the LU objective to a small-step reverse KL and establishes how per-token times preserve the forward marginal and support confidence-based sampling.

A.2.2 Proof of Theorem 1

  • As the reverse step shrinks, the centered reverse KL retains a finite model-dependent part; weighting it yields the LU objective on corrupted positions.The added logarithm is independent of model parameters and removes the divergent part of the small-step KL.Source: §A.2.2 · p. 16

A.3 Per-Token Time Embeddings

  • For a linear noise schedule, averaging over per-token times recovers the original uniform-state marginal, and independent times preserve the factorized sequence law.For a nonlinear schedule, the proof requires E_qt[τ_i] = 1 − α_t.Source: §A.3 · p. 17
  • During denoising, committed positions receive τ_i = 0 and become clean conditioning context, while positive-time positions remain targets for denoising.Source: §A.3 · p. 17
  • The per-token-time design yields the confidence-based easy-to-hard sampler used by LUDI.Source: §A.3 · p. 17

B.1 Training and Sampling

LUDI trains with token-level corruption hints and a loss that encourages reverse transitions toward clean tokens, then samples by committing high-confidence tokens. The AR-to-block-diffusion adaptation adjusts prediction positions, attention, and auxiliary training to retain useful AR behavior.

B.1 Training and Sampling

  • LUDI combines per-token time embeddings with the LU loss and a confidence-based scheduler that progressively commits high-confidence tokens as conditions for few-step generation.Source: §B.1 · p. 17

B.2 Details of AR to Block Diffusion

  • Label shifting resolves the mismatch between AR next-token prediction and block diffusion’s prediction positions while preserving diffusion and AR prediction preference.The shift changes the prediction position and still allows token ℓ−1 to attend to token ℓ.Source: §B.2 · p. 18
  • Complementary mask views make each position serve as both a clean target and a noisy input in one forward pass, providing doubled supervision in the low-data regime.Source: §B.2 · p. 18
  • The context-causal mask keeps committed context causal while allowing bidirectional attention within the active denoising block, enabling KV-cache reuse and parallel token decoding at inference.Figure 5 contrasts the training and inference masks; at inference, the prompt and past blocks form the frozen causal prefix.Source: Figure 5 · p. 19, §B.2 · p. 19
  • An auxiliary autoregressive loss on the clean sequence is computed without extra forward cost and is intended to prevent catastrophic forgetting while stabilizing early training.The objective is L = L_diff + λL_AR.Source: §B.2 · p. 19

C.1 Experimental Setup for Small-Scale LUDI

The experimental setup describes 170M UDLM training and 1B UDLM training from scratch, with a table caption identifying a mathematical-reasoning evaluation.

170M models

  • Four 170M UDLMs using LUDI and three loss baselines were trained on OpenWebText for 500K iterations, with generation metrics including generative perplexity and token-level entropy.The setup specifies a 12-block Transformer and compares against released RADD and MDLM checkpoints trained on the same dataset.Source: §C.1 · p. 19

1B models

  • Four 1B-parameter UDLMs were trained from scratch on FineWeb for 500K iterations and evaluated on six multiple-choice benchmarks.Accuracy is reported for PIQA, SIQA, ARC-Easy, and RACE; length-normalized accuracy is reported for HellaSwag and OpenBookQA.Source: §C.1 · p. 19
  • Table 4 is identified as reporting mathematical-reasoning accuracy with the Dolci and Dolci-Nemo mixture.Source: Table 4 · p. 20

C.2 Experimental Setup for LUDI-7B

The evaluation compares LUDI-7B with similarly scaled autoregressive and diffusion models across code, mathematical reasoning, instruction following, and knowledge-intensive question answering.

C.2 Experimental Setup for LUDI-7B

  • The baselines include Qwen2.5-7B-Instruct as the AR initialization and masked diffusion models ranging from 7B to 8B.The masked diffusion baselines are Dream 7B, LLaDA 8B, LLaDA-1.5 8B, LLaDA-MoE 7B, and Fast-dLLM v2.Source: §C.2 · p. 20
  • Evaluation spans code generation, mathematical reasoning, instruction following, and knowledge-intensive question answering.The benchmarks include HumanEval, HumanEval+, MBPP, MBPP+, GSM8K, IFEval, MMLU, and GPQA.Source: §C.2 · p. 20
  • Code benchmarks use EvalPlus, and all tasks report standard accuracy following the Fast-dLLM v2 evaluation protocol.Source: §C.2 · p. 20

C.3 Mathematical Reasoning

The appendices analyze mathematical reasoning and inference efficiency, then examine the LU loss implementation, training convergence, and sample quality and diversity.

C.3 Mathematical Reasoning

  • LUDI improves from 75.2 to 82.1 on GSM8K and 49.0 to 52.6 on MATH500 with Nemotron-Post-Train v2, reaching the strongest compared MATH500 result.Mathematical scores are substantially lower under the Dolci setting across UDLM and MDLM variants; the authors attribute this shared pattern primarily to training–evaluation distribution mismatch.Source: §C.3 · p. 20
  • Repetitive loops make up 21.5% of LUDI’s incorrect MATH500 responses, versus 27.8% for Fast-dLLM v2.The authors describe this as a moderate reduction in a diffusion-specific failure mode and note that generated tokens can be revised.Source: §C.3 · p. 21

C.4 End-to-End Inference Efficiency

  • At batch sizes 1, 4, and 16, LUDI throughput speedups are 1.92×, 1.30×, and 1.03×; AR slightly leads at batch size 32.LUDI uses approximately 1.9–2.0× the baseline peak memory and has higher end-to-end latency, reflecting longer reasoning chains; Table 5 compares the systems on GSM8K.Source: §C.4 · p. 21, Table 5 · p. 21
  • The Dolci mathematical subset has a median response length of 2,114 tokens, compared with 217 tokens for GSM8K, contributing to LUDI’s higher end-to-end latency.Source: §C.4 · p. 21

C.5 Efficiency of the LU Loss

  • The optimized LU operator achieves a 4.39× speedup and 3.45× memory reduction over the PyTorch ELBO baseline.It avoids vocabulary-wide nonlinear operations after normalization by applying them only at indices m and y, reducing intermediate tensors, operations, and kernel launches; Figure 6 covers a long-sequence, large-vocabulary setting.Source: §C.5 · p. 22, Figure 6 · p. 22

C.6 Convergence

  • At 500K steps, LUDI has the lowest generative perplexity under both 512-step and 1024-step sampling, following rapid early improvement.Figure 7 compares training trajectories for four UDLM objectives under the same evaluation protocol.Source: §C.6 · p. 22, Figure 7 · p. 22

C.7 Quality and Diversity

  • At the 500K checkpoint, LUDI’s mean Gen.PPL is 37.9 and mean entropy is 7.50, without an obvious entropy collapse.Figure 8 shows per-sample Gen.PPL and entropy distributions from 1024-step sampling with 1024 samples per loss; each panel has its own axis and density scale.Source: §C.7 · p. 23, Figure 8 · p. 23

Related papers

Loading 2609.35817v1…