Consistent Diffusion Language Models

1. Problem Statement

Diffusion Language Models (DLMs) offer a promising alternative to autoregressive (AR) models by enabling parallel token generation and sublinear-time inference. However, a critical bottleneck remains: generating high-quality text typically requires hundreds of refinement steps, negating the efficiency gains promised by parallelism.

In continuous domains (e.g., images), Consistency Models have successfully accelerated diffusion by training models to be invariant along a unique, deterministic trajectory defined by the Probability Flow Ordinary Differential Equation (PF-ODE).

• The Gap: In discrete domains (like text), no analogous sample-space PF-ODE exists. Discrete corruption processes (e.g., masking or uniform noise) do not follow a single deterministic path; instead, they involve stochastic transitions.
• The Consequence: Directly adapting continuous consistency methods to discrete data is ill-defined. Existing solutions often rely on multi-stage distillation pipelines (using a teacher model) or continuous-relaxation surrogates, which are complex, computationally expensive, and often suffer from mode collapse or diversity loss.

2. Methodology: Multi-Path Discrete Consistency (MPDC)

The authors propose a fundamental shift: instead of searching for a non-existent deterministic path, they leverage the exact posterior bridges inherent in discrete diffusion processes.

Core Concept: The Posterior Bridge

For any two noise levels $s < t$ in a discrete diffusion process, there exists a closed-form exact posterior bridge $q(x_s | x_t, x_0)$. This bridge defines a valid stochastic path connecting a noisy state $x_t$ to a cleaner state $x_s$, conditioned on the true data $x_0$.

• Unlike the continuous PF-ODE, this bridge is stochastic.
• Crucially, these bridges form a semigroup (Chapman-Kolmogorov property), meaning paths can be composed transitively.

The MPDC Principle

The authors introduce Multi-Path Discrete Consistency (MPDC), which trains a denoiser to be path-invariant in expectation.

• Definition: A model $f_\theta(x_t, t)$ is consistent if its prediction at a noisy state $x_t$ agrees (in expectation) with the prediction made after taking an intermediate "hop" to a cleaner state $x_s$ via the exact posterior bridge.
• Objective: The model minimizes the divergence between its direct prediction from $x_t$ and the expected prediction from $x_s$ (where $x_s$ is sampled from the bridge).
• Mathematical Formulation:
  $$\mathcal{L}_{CDLM} = \mathbb{E}_{t, \delta, x_t} \left[ \sum_{i \in \mathcal{M}(x_t)} w(t, \delta) \cdot D\left( f_\theta(x_t, t)_i \parallel \tilde{f}_\theta(x_s, s)_i \right) \right]$$
  Where $x_s \sim q(x_s | x_t, x_0)$, $\tilde{f}_\theta$ is a target network (EMA), and $D$ is a divergence metric (typically Jensen-Shannon Divergence).

Key Design Choices for Stability

To prevent mode collapse (where the model outputs repetitive text to trivially satisfy consistency) and uniform drift, the authors introduce three stabilizing mechanisms:

1. Step Size as Curriculum: Sampling step sizes $\delta$ from a specific range (e.g., $1/8$ to $3/8$) and weighting them by $1/\delta$ to balance learning signals across short and long paths.
2. Max-Step Diffusion Anchor: A small fraction of the training loss is anchored to the true data ($s=0$) using a standard diffusion objective. This prevents the model from converging to degenerate constant solutions.
3. Optimization Asymmetry & Divergence: Using a stop-gradient on the target network and employing Jensen-Shannon Divergence (JSD) instead of Forward/Backward KL. JSD provides a stable trade-off between mode-seeking and mode-covering, avoiding the instability of KL divergence when training from scratch.

3. Key Contributions

1. A New Principle for Discrete Generative Modeling: The paper establishes that exact posterior bridges serve as a rigorous, closed-form substitute for the missing PF-ODE, turning the "discrete acceleration" obstacle into an analytic asset.
2. A Unified Training Framework: CDLM provides a single-stage, teacher-free objective. The authors prove that standard masked diffusion, continuous consistency models, progressive distillation, and recent two-stage distillation methods (like SDTT and DUO-DCD) are all special cases, limits, or approximate couplings of the CDLM framework.
3. State-of-the-Art Performance: CDLM achieves superior results in both unconditional and conditional text generation, outperforming strong base models and matching or surpassing complex multi-stage distilled models without requiring a teacher.

4. Experimental Results

The authors evaluated CDLM (specifically MCDLM for masked diffusion and UCDLM for uniform diffusion) against baselines like MDLM, DUO, SDTT, and DUO-DCD on the OpenWebText dataset.

• Unconditional Generation:

  • Few-Step Efficiency: CDLM establishes new state-of-the-art (SOTA) results for single-stage models. It significantly outperforms base models (MDLM, DUO) across all sampling budgets.
  • Comparison to Distillation: Remarkably, the single-stage CDLM often outperforms multi-stage distilled models (SDTT, DUO-DCD) in terms of Perplexity (PPL) while maintaining much higher entropy (diversity).
  • Mode Collapse: Distilled models using greedy samplers (e.g., DUO-DCD) showed low entropy (3.9), indicating severe mode collapse (repetitive characters). CDLM maintained healthy entropy (5.0+), proving its ability to preserve diversity.
  • Speed: CDLM achieves 32x–64x speedups over autoregressive baselines and 64x–128x over standard MDLM while maintaining comparable quality.

• Conditional Generation:

  • CDLM demonstrated robust performance in recovering original sentences from 50% corrupted inputs across multiple datasets (OpenWebText, Lambada, Wikitext-103, PTB).
  • It achieved better BLEU scores and lower perplexity than SDTT and DUO, particularly in preserving input conditions where other models failed.

5. Significance and Impact

• Theoretical Unification: The paper provides a "unifying lens" for discrete generative modeling, showing that diverse acceleration techniques are fundamentally connected through the concept of path consistency.
• Scalability: By eliminating the need for multi-stage distillation and teacher models, CDLM offers a principled, scalable, and teacher-free foundation for next-generation fast discrete generative models.
• Generalizability: While demonstrated on text, the framework applies to any discrete domain with a tractable posterior bridge (e.g., biological sequences, graphs, program synthesis).
• Paradigm Shift: It reframes fast discrete diffusion not as an approximation of continuous ODEs, but as the learning of a path-independent denoiser over stochastic bridges, successfully solving the few-step generation problem in discrete spaces.

In conclusion, CDLM bridges the gap between the efficiency of continuous consistency models and the realities of discrete data, offering a robust, single-stage solution that achieves high-fidelity, fast text generation without the complexity of distillation pipelines.