Table of Contents
Fetching ...

Continuous Diffusion Model for Language Modeling

Jaehyeong Jo, Sung Ju Hwang

TL;DR

This work introduces RDLM, a continuous diffusion model for language and discrete data that leverages the geometry of the categorical probability manifold. By mapping discrete tokens to continuous states on the hypersphere and formulating diffusion as a flow on this manifold, the method generalizes discrete diffusion, enables simulation-free training via radial symmetry, and employs a mixture-path strategy to balance diffusion behaviors. The approach yields strong performance on language benchmarks (e.g., Text8, LM1B), and demonstrates applicability to image and DNA sequence design, surpassing several discrete diffusion baselines and approaching autoregressive models. The combination of geometry-aware reparameterization, tractable training objectives, and dimension-splitting makes RDLM scalable to large vocabularies while preserving the benefits of iterative refinement.

Abstract

Diffusion models have emerged as a promising alternative to autoregressive models in modeling discrete categorical data. However, diffusion models that directly work on discrete data space fail to fully exploit the power of iterative refinement, as the signals are lost during transitions between discrete states. Existing continuous diffusion models for discrete data underperform compared to discrete methods, and the lack of a clear connection between the two approaches hinders the development of effective diffusion models for discrete data. In this work, we propose a continuous diffusion model for language modeling that incorporates the geometry of the underlying categorical distribution. We establish a connection between the discrete diffusion and continuous flow on the statistical manifold, and building on this analogy, introduce a simple diffusion process that generalizes existing discrete diffusion models. We further propose a simulation-free training framework based on radial symmetry, along with a simple technique to address the high dimensionality of the manifold. Comprehensive experiments on language modeling benchmarks and other modalities show that our method outperforms existing discrete diffusion models and approaches the performance of autoregressive models. The code is available at https://github.com/harryjo97/RDLM.

Continuous Diffusion Model for Language Modeling

TL;DR

This work introduces RDLM, a continuous diffusion model for language and discrete data that leverages the geometry of the categorical probability manifold. By mapping discrete tokens to continuous states on the hypersphere and formulating diffusion as a flow on this manifold, the method generalizes discrete diffusion, enables simulation-free training via radial symmetry, and employs a mixture-path strategy to balance diffusion behaviors. The approach yields strong performance on language benchmarks (e.g., Text8, LM1B), and demonstrates applicability to image and DNA sequence design, surpassing several discrete diffusion baselines and approaching autoregressive models. The combination of geometry-aware reparameterization, tractable training objectives, and dimension-splitting makes RDLM scalable to large vocabularies while preserving the benefits of iterative refinement.

Abstract

Diffusion models have emerged as a promising alternative to autoregressive models in modeling discrete categorical data. However, diffusion models that directly work on discrete data space fail to fully exploit the power of iterative refinement, as the signals are lost during transitions between discrete states. Existing continuous diffusion models for discrete data underperform compared to discrete methods, and the lack of a clear connection between the two approaches hinders the development of effective diffusion models for discrete data. In this work, we propose a continuous diffusion model for language modeling that incorporates the geometry of the underlying categorical distribution. We establish a connection between the discrete diffusion and continuous flow on the statistical manifold, and building on this analogy, introduce a simple diffusion process that generalizes existing discrete diffusion models. We further propose a simulation-free training framework based on radial symmetry, along with a simple technique to address the high dimensionality of the manifold. Comprehensive experiments on language modeling benchmarks and other modalities show that our method outperforms existing discrete diffusion models and approaches the performance of autoregressive models. The code is available at https://github.com/harryjo97/RDLM.

Paper Structure

This paper contains 83 sections, 11 theorems, 86 equations, 6 figures, 6 tables, 3 algorithms.

Key Result

Proposition 3.1

The transition distribution of discrete diffusion processes can be modeled by the continuous flow on the statistical manifold, and further on the hypersphere.

Figures (6)

  • Figure 1: Illustration of the continuous reparameterization of discrete data and two types of our generative process on hypersphere. (a) Example of a transition distribution of a discrete diffusion process modeled by a continuous flow on a $d$-dimensional sphere. (b) Illustration of the diffusion processes on $\mathbb{S}^2$ generalizing masked diffusion and uniform diffusion, respectively.
  • Figure 2: Maximum mean discrepancy (MMD) distance between the simulated distribution $p(\bm{X}_t|\bm{X}_0,\bm{X}_T)$ and the approximated distribution. We report the results for dimensions 4, 256, and 30522.
  • Figure 5: Comparison between the training objectives. We compare Bits Per Character (BPC) on the Text8 test set.
  • Figure : Training
  • Figure : Sampling
  • ...and 1 more figures

Theorems & Definitions (15)

  • Proposition 3.1
  • Lemma A.1
  • proof
  • Proposition A.2
  • proof
  • Corollary A.3
  • Proposition A.4
  • proof
  • Corollary A.5
  • Proposition A.6
  • ...and 5 more