Towards Latent Diffusion Suitable For Text

TL;DR

The paper tackles efficient language generation with diffusion models by learning a data-conditioned forward diffusion in latent space. It introduces Neural Flow Diffusion Models (NFDM), extending prior NFDM frameworks to language by modeling $q_{\varphi}(\mathbf{z}_t|\mathbf{x})$ via a forward transformation $F_{\varphi}$ and learning a reverse process, with forward SDE $\mathrm{d}\mathbf{z}_t = \tilde{f}^{F}(\mathbf{z}_t,t,\mathbf{x})\,dt + g_{\varphi}^{2}(t)\,\mathrm{d}\mathbf{w}$ and reverse $\mathrm{d}\mathbf{z}_t = \tilde{f}^{B}(\mathbf{z}_t,t,\mathbf{x})\,dt + g(t)\,\mathrm{d}\bar{\mathbf{w}}$. They show that NFDM reduces the likelihood gap with autoregressive models of similar size and achieves sample quality competitive with prior diffusion baselines on unconditional generation tasks. The work highlights the benefits of adaptive, per-token noise schedules and data-conditioned forward processes, while noting limitations such as dataset scale and lack of conditional generation experiments. Overall, NFDM demonstrates that learning a forward diffusion conditioned on input can tighten likelihood objectives and offer robust generation, motivating further exploration on conditional tasks and larger corpora.

Abstract

Language diffusion models aim to improve sampling speed and coherence over autoregressive LLMs. We introduce Neural Flow Diffusion Models for language generation, an extension of NFDM that enables the straightforward application of continuous diffusion models to discrete state spaces. NFDM learns a multivariate forward process from the data, ensuring that the forward process and generative trajectory are a good fit for language modeling. Our model substantially reduces the likelihood gap with autoregressive models of the same size, while achieving sample quality comparable to that of previous latent diffusion models.

Figures (8)

• Figure 1: Diffusion-LM $\gamma(t)$ for $t \in [0,1]$
• Figure 2: Diffusion-LM $\dot\gamma(t)$ for $t \in [0,1]$
• Figure 3: Diffusion-LM $g(t)^2$ for $t \in [0,1]$
• Figure 4: Cosine‐similarity between $\overline{\mu}_{\varphi}(E_\varphi({\bm{x}}), t)$ and $\overline{\mu}_{\varphi}(E_\varphi({\bm{x}}), t+\Delta t)$, with steps of $t=0.1$ comparing NFDM-Additive (left) and NFDM (right). Standard deviation computed within batch.
• Figure 5: Samples generated using discrete sampling with $\tilde{\sigma}_{s|t}=1$ and 2000 steps.