SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Justin Lovelace* Varsha Kishore Chao Wan Eliot Shekhtman Kilian Q. Weinberger
Cornell University, Ithaca, NY
Diffusion models have achieved great success in modeling continuous data modalities such as images, audio, and video, but have seen limited use in discrete domains such as language. Recent attempts to adapt diffusion to language have presented diffusion as an alternative to existing pretrained language models. We view diffusion and existing language models as complementary. We demonstrate that encoder-decoder language models can be utilized to efficiently learn high-quality language autoencoders. We then demonstrate that continuous diffusion models can be learned in the latent space of the language autoencoder, enabling us to sample continuous latent representations that can be decoded into natural language with the pretrained decoder. We validate the effectiveness of our approach for unconditional, class-conditional, and sequence-to-sequence language generation. We demonstrate across multiple diverse data sets that our latent language diffusion models are significantly more effective than previous diffusion language models. Our code is available at https://github.com/justinlovelace/latent-diffusion-for-language.
Although originally introduced by Sohl-Dickstein et al. [61] in 2015, diffusion models did not see widespread use until Ho et al. [22] demonstrated their viability for high-quality image generation in 2020. Since then, research has driven rapid improvements and they have recently surpassed generative adversarial networks on image generation benchmarks [12] and autoregressive models on density estimation benchmarks [30], outclassing generative modeling paradigms that have dominated those areas for the better part of a decade. Diffusion models are now, arguably, the most widely used class of generative models for continuous data modalities such as images, audio, and video [54, 32, 23].
The widespread success of diffusion models across a variety of domains and applications makes them appealing for language generation. However, they have seen less use in discrete domains, where the gradual transition of discrete states to Gaussian noise (and vice versa) is not as natural as in continuous domains. Prior work proposes to learn continuous diffusion models in the space of learnable word embeddings and decodes the continuous generations with a rounding step [36]. However, combining representation learning with the diffusion objective requires careful regularization to avoid collapse.
One breakthrough in image generation was the introduction of latent diffusion [51], where diffusion models are trained to produce samples from the latent distribution of a pretrained autoencoder. This offloads the task of generating high-frequency details to the autoencoder and enables the diffusion process to focus on the high-level semantics of images. In this paper, we explore the viability of latent diffusion for text generation. We claim that this approach is particularly well-suited for discrete modalities because it offloads the challenge of modeling a discrete distribution to the autoencoder and simplifies the diffusion process by restricting it to the continuous, latent feature space.
We introduce Latent Diffusion for Language Generation (LD4LG), a method that leverages the latent space of a pretrained encoder-decoder network (e.g. BART [35], T5 [50]) to learn a high-quality
*Correspondence to <jl3353@cornell.edu>.
diffusion model for text. The latent representations from such models are high-dimensional and input-length dependent — complicating the use of diffusion models [51, 66]. To address both issues, we learn an additional compression module that maps the high-dimensional encoder representations to a lower-dimensional fixed-length representation. We also learn a corresponding reconstruction network to map these fixed-length features back to high dimensional features that guide the language decoder (via cross-attention) to reconstruct the original language.
The low-dimensional representation is ideally suited for diffusion. For language generation, we use a diffusion model to generate a low-dimensional (fixed-length) latent, which is mapped into a higher dimensional space with the reconstruction network. This high dimensional representation then guides the pre-trained decoder to generate natural language. Our approach naturally combines the continuous, fixed-length diffusion process with discrete, variable length text generation.
We demonstrate that LD4LG is effective for unconditional, class-conditional, and sequence-to-sequence language generation across a variety of datasets. Our approach significantly outperforms recent diffusion language models while using fewer sampling steps. For instance, we achieve a MAUVE score [47] of .716 for the ROCStories dataset with 250 sampling steps while DiffusionLM [36] achieves a MAUVE score of .043 using 2000 sampling timesteps. For the challenging XSum summarization benchmark, we achieve a ROUGE-L of 31.9 with 250 timesteps while the recently proposed DiffuSeq [16] achieves a ROUGE-L of 14.1 with 2000 timesteps. We also find that the diffusion models offer some benefits over a strong autoregressive baseline. In particular, we observe that our latent language diffusion is less susceptible to memorization and more effective for class-conditional generation.
Diffusion models [61, 22, 63] are a class of latent variable models that learn to iteratively transform random Gaussian noise, which can be sampled analytically, to a sample from an unknown data distribution specified by a collection of samples. This mapping is defined through a forward diffusion process that iteratively adds Gaussian noise to samples, and a generative process that iteratively "denoises" samples from the Gaussian distribution to obtain samples from the data distribution. We provide a formal description of diffusion models in the appendix.
The diffusion model consists of a denoising network  $\hat{\mathbf{x}}_{\theta}$  trained with a regression objective
$$
\mathcal {L} (\theta) = \mathbb {E} _ {t, \mathbf {x}, \epsilon} [ \lambda_ {t} \| \hat {\mathbf {x}} _ {\theta} (\sqrt {\alpha_ {t}} \mathbf {x} + \sqrt {1 - \alpha_ {t}} \epsilon , t) - \mathbf {x} \| _ {2} ^ {2} ]
$$
where  $\mathbf{x}$  is the training data,  $t\sim \mathcal{U}(\mathbf{0},\mathbf{1})$  is the timestep,  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{1})$  is Gaussian noise,  $\alpha_{t}$  defines the noise schedule, and  $\lambda_{t}$  is a time-dependent weighting term. The denoising network is therefore trained to denoise a noisy latent,  $\mathbf{z}_t = \sqrt{\alpha_t\mathbf{x}} +\sqrt{1 - \alpha_t\epsilon}$ , to the clean data,  $\mathbf{x}$ , with a regression objective that emphasizes certain times  $t$ . Sampling algorithms start from pure Gaussian noise,  $\mathbf{z}_1\sim \mathcal{N}(\mathbf{0},\mathbf{1})$ , and utilize the denoising network to iteratively generate latents  $\mathbf{z}_{t_1},\mathbf{z}_{t_2},\dots,\mathbf{z}_{t_T}$  where  $1 = t_{1} > t_{2} > \ldots >t_{T} = 0$ , with decreasing levels of noise until  $\mathbf{z}_0$  is drawn approximately from the data distribution.
3 Latent Diffusion For Language
Figure 1 presents an overview of Latent Diffusion for Language Generation. Our method consists of two main parts. We augment a pretrained encoder-decoder language model with two learnable networks to develop a high-quality language autoencoder with a compact latent space. We then introduce continuous diffusion models that learn to generate samples from the latent distribution of our language autoencoders. These continuous samples can, by design, be decoded into natural language.
3.1 Language Autoencoder
We base our architecture on pretrained encoder-decoder language models (depicted in blue), such as BART [35] and T5 [50] (we present results with both). By default, we freeze the pre-trained models and learn only the autoencoding modules to accelerate training. The Language Encoder,  $E(\cdot)$ , maps variable-length language, represented as a sequence of tokens,  $\mathbf{w} \in \mathbb{N}^L$ , to a latent representation of the same length,  $E(\mathbf{w}) \in \mathbb{R}^{L \times d_{\mathrm{LM}}}$ .
Figure 1: Overview of our proposed latent language diffusion framework.
Compression Network. The learnable Compression Network maps the encoder features to a compact latent space that is well-suited for diffusion. We adopt the Perceiver Resampler [2] architecture, originally developed to compress image features for a vision-language model, which is depicted in Figure 2. The Perceiver Resampler, like the transformer, consists of a stack of alternating multi-head attention (MHA) blocks and feedforward (FF) layers. We refer the reader to Vaswani et al. [69] for a detailed description of these components. We learn  $\ell$  latent queries  $Z\in \mathbb{R}^{\ell \times d_{\mathrm{LM}}}$  that iteratively cross-attend to the language encoder features  $E(\mathbf{w})\in \mathbb{R}^{L\times d_{\mathrm{LM}}}$  to extract information. We follow Alayrac et al. [2] and allow the latent queries to simultaneously attend to themselves and the frozen encoder representations. We can write the attention layer as
$$
Z = Z + \operatorname {M H A} (\mathrm {q} = Z, \mathrm {k v} = [ Z; E (\mathrm {w}) ])
$$
where  $\mathrm{MHA}(\cdot)$  is the multi-head attention operation with queries, q, and keys/values, kv. This design compresses the encoder representations to the fixed sequence length,  $\ell$ , of the latents. After each multi-head attention layer, a feedforward layer is applied to the latent query representations.
After the Compression Network maps the input to a fixed sequence length, we reduce the dimensionality of the output to dimension  $d_{\mathrm{ae}}$  with a learnable linear projection. The compression network therefore maps the variable length output of the frozen encoder to a compact latent space
$$
\mathbf {x} = f _ {\phi} (E (\mathbf {w})) \in \mathbb {R} ^ {\ell \times d _ {\mathrm {a c}}}
$$
of fixed length  $\ell < L$  and dimensionality  $d_{\mathrm{ae}} < d_{\mathrm{LM}}$  where we will learn our diffusion model.
To ensure that the latent space is appropriately scaled for diffusion, we can optionally constrain the norm of the latent space. Since  $\mathbb{E}_{\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I})}[\| \epsilon \|_2^2] = d_{\mathrm{ae}}$  [1],
we can normalize the latent vectors along the feature dimension so that  $\| \mathbf{x}_i\| _2^2 = d_{\mathrm{ae}}$  similar to prior work on text diffusion [13].
Figure 2: Architecture of our Compression Network.
Reconstruction Network. The Reconstruction Network maps the compressed latent space to the feature space expected by the Language Decoder. To achieve this, we project  $\mathbf{x} = f_{\phi}(E(\mathbf{w})) \in \mathbb{R}^{\ell \times d_{\mathrm{ae}}}$  back up to dimension  $d_{\mathrm{LM}}$ , add learnable absolute position embeddings, and pass it through a standard transformer model to obtain features  $g_{\phi}(\mathbf{x}) \in \mathbb{R}^{\ell \times d_{\mathrm{LM}}}$ .
The Language Decoder,  $D(\cdot)$ , cross-attends to to these features and generates text autoregressively. We train the compression and reconstruction networks to produce features that guide the decoder to reconstruct the input text
$$
\mathbf {w} \approx D (g _ {\phi} (\mathbf {x})) = D (g _ {\phi} (f _ {\phi} (E (\mathbf {w})))
$$
with the cross-entropy loss. This gives us a continuous, semantic latent space that can be decoded to natural language.
Implementation Details. We utilize BART-base and FLAN-T5-base [10] as the encoder-decoder language models throughout this work and learn language autoencoders for each dataset. During autoencoder training, we freeze the pre-trained language models and only learn the autoencoding modules. The autoencoder training could likely be amortized across datasets by training a general-purpose language autoencoder on a large corpus of text, but we leave such explorations to future work. We train the autoencoder to reconstruct the input language with the cross-entropy loss. For the diffusion latent space, we set  $\ell = 32$ ,  $d_{\mathrm{ae}} = 64$  and utilize 3 layers in both autoencoding modules across all monolingual datasets.
For our machine translation experiments, we utilize MT5-base [72] to develop our autoencoder. We found it beneficial to jointly fine-tune the language model and the autoencoding modules, likely because the dataset is an order of magnitude larger than our other datasets and therefore benefits from the additional capacity. We use the same latent dimensionality, but only use a single layer for the autoencoding modules. We report full hyperparameter settings in the appendix. We constrain the norm of the latent space across models and datasets except when using FLAN-T5 because it led to a minor degradation in autoencoding performance and downstream generation quality.
3.2 Latent Language Diffusion
Figure 1 outlines our latent language diffusion framework. Given some dataset of natural language,  $\mathcal{D}$ , we can now sample continuous data as  $\mathbf{x} = f_{\phi}(E(\mathbf{w})) \in \mathbb{R}^{\ell \times d_{\mathrm{ae}}}$  where  $\mathbf{w} \sim \mathcal{D}$ . We then train a continuous denoising network,  $\hat{\mathbf{x}}_{\theta}(\cdot)$ , to recover  $\mathbf{x}$  with the standard regression objective
with some time-dependent weighting  $\lambda_{t}$ . In practice, the denoising network is often parameterized as an  $\epsilon$ -prediction network [22] or a  $\mathbf{v}$ -prediction network [57] where the velocity,  $\mathbf{v}$ , is defined as  $\mathbf{v} = \sqrt{\alpha_t\epsilon} - \sqrt{1 - \alpha_t\mathbf{x}}$ . These parameterizations can be interpreted as different weighting functions,  $\lambda_{t}$ , for the regression objective above (see Salimans and Ho [57]. We adopt the  $\mathbf{v}$ -parameterization in this work because it has been shown to be effective for latent image diffusion [51].
For generation, we sample a latent variable,  $\mathbf{z}_1\in \mathbb{R}^{\ell \times d_{\mathrm{ac}}}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ , that is iteratively denoised to produce a sample,  $\mathbf{x} = \mathbf{z}_0$ , from the distribution of the language autoencoder's latent space. We then generate natural language with the pretrained reconstruction network and language decoder  $\mathbf{w} = D(g_{\phi}(\mathbf{x}))$ . We train our diffusion models with the cosine noise schedule  $\alpha_{t} = \cos (0.5\pi t)^{2}$  [45, 57, 55] by default. For our machine translation experiments, we employ a scaled cosine noise schedule (see subsection E.2 in the appendix for full details) [7, 27]. For generation, we use the DDPM sampler with 250 sampling timesteps. For text generation with the pretrained decoder, we utilize beam search with 4 beams. We train all of our diffusion models with a single Nvidia A6000 GPU except for the machine translation models which are trained with 4 Nvidia A6000 GPUs.
Denoising Network Architecture. Our denoising network,  $\hat{\mathbf{x}}_{\theta}(\mathbf{z}_t,t)$ , is a pre-LayerNorm transformer [69, 70] with 12 layers and a dimension of 768. We utilize learnable absolute positional encodings and GeGLU activations [59]. Bao et al. [4] adapted transformers to image diffusion and found that dense connections [28] between early and late layers are beneficial due to the dense nature of the denoising objective. We adopt this modification to improve the suitability of the transformer for diffusion. The autoencoder latent is projected to the transformer dimension, processed by the transformer, and then projected back to dimensionality of the autoencoder latent to obtain the final prediction. Following prior work [6, 56, 7], we utilize  $\alpha$ -conditioning to condition the model on the level of noise. We map  $\alpha_{t}$  to a sinusoidal positional embedding [69] and pass it through an MLP with a single hidden layer to obtain a time embedding. We add this time embedding to the input sequence and apply adaptive layer normalization [46] conditioned on the time embedding to the output of every feedforward layer.
Self-Conditioning We utilize the self-conditioning technique introduced by Chen et al. [8] which has been shown to improve the quality of diffusion models [8, 67]. The denoising network is typically conditioned on the latent variable and the current timestep as  $\tilde{\mathbf{x}}_t = \hat{\mathbf{x}}_\theta (\mathbf{z}_t,t)$ . Self-conditioning proposes to condition the network on its estimate of the data from the previous timestep,  $s > t$ , to improve the prediction at the current timestep  $\tilde{\mathbf{x}}_t = \hat{\mathbf{x}}_\theta (\mathbf{z}_t,t,\tilde{\mathbf{x}}_s)$ . During inference, the sampling procedure is inherently iterative and at time  $t$ , we have already computed the output of the denoising network for the previous step. Therefore, it does not require any additional applications of the
network. We must, however, modify the training procedure so that the denoising network learns to utilize the estimate of the data, and we must define the inference behavior for the first timestep.
For each training step, we sample some time  $t \sim \mathcal{U}([0,1])$  as before. With probability  $p$ , we do not provide any estimate of the data for self-conditioning, denoted  $\tilde{\mathbf{x}}_{t,\emptyset} = \hat{\mathbf{x}}_{\theta}(\mathbf{z}_{t},t,\emptyset)$ . With probability  $1 - p$ , however, we mimic the inference behavior by first computing  $\tilde{\mathbf{x}}_{t,\emptyset} = \hat{\mathbf{x}}_{\theta}(\mathbf{z}_{t},t,\emptyset)$  and then computing an additional estimate  $\tilde{\mathbf{x}}_t = \hat{\mathbf{x}}_\theta (\mathbf{z}_t,t,\mathrm{sg}(\tilde{\mathbf{x}}_{t,\emptyset}))$  where  $\mathrm{sg}()$  is the stop-gradient operation. This second estimate is then used to compute the loss. We follow Chen et al. [8] and set  $p = 0.5$ .
This training procedure also maintains the capacity for inference without self-conditioning which is utilized to generate the first estimate during sampling. We condition on the previous estimate by concatenating it with the noisy latent along the feature dimension. When the previous estimate is not provided, we concatenate a learnable embedding with the noisy latent.
Class-Conditional Diffusion. For class-conditional diffusion, we have some dataset where each natural language utterance is associated with one of  $C$  class labels representing, for example, the topic of the text. We condition the denoising network on the class label,  $y$ , during training,  $\hat{\mathbf{x}}_t = \hat{\mathbf{x}}_\theta(\mathbf{z}_t, t, y)$ . We replace the ground truth class label,  $y_i$ , with a null label,  $y_\emptyset$ , with probability  $p = 0.1$  to maintain the capacity for unconditional generation. At inference time, we can choose some class  $y$  to guide the sampling process to generate text from the specified class. We condition on class labels by introducing learnable embeddings for all labels, including the null label, and add it to the time embedding.
Sequence-to-Sequence Diffusion. Given some seq2seq dataset consisting of source-target language pairs  $(\mathbf{w}_{\mathrm{src}},\mathbf{w}_{\mathrm{trg}})\sim \mathcal{D}$ , we condition our denoising network on the source sequence and generate the target latent  $\mathbf{x}_{\mathrm{trg}} = f_{\phi}(E(\mathbf{w}_{\mathrm{trg}}))$ . For news summarization, for instance, we generate a latent representation of the summary by conditioning the network on the article text. To achieve this, we introduce a cross-attention layer after every self-attention layer in the denoising network that attends to features from a frozen language encoder.
In general, we can incorporate any language encoder,  $E_{\mathrm{src}}(\cdot)$ , to extract features from the source text. By default, we use the same pretrained encoder used for our language autoencoder. For our machine translation experiments, we condition our latent diffusion models on representations from a frozen MT5-XL encoder, which we found to be more effective than MT5-base representations. Therefore, given a sample from our seq2seq dataset,  $(\mathbf{w}_{\mathrm{src}}, \mathbf{w}_{\mathrm{trg}}) \sim \mathcal{D}$ , we can compute  $\mathbf{x}_{\mathrm{trg}} = f_{\phi}(E(\mathbf{w}_{\mathrm{trg}}))$  and use a modified seq2seq diffusion objective
$$
\mathcal {L} (\theta) = \mathbb {E} _ {t, \left(\mathbf {w} _ {\mathrm {s r c}}, \mathbf {w} _ {\mathrm {t r g}}\right), \epsilon} [ \lambda_ {t} \| \hat {\mathbf {x}} _ {\theta} (\sqrt {\alpha_ {t}} \mathbf {x} _ {\mathrm {t r g}} + \sqrt {1 - \alpha_ {t}} \epsilon , t, E _ {\mathrm {s r c}} (\mathbf {w} _ {\mathrm {s r c}})) - \mathbf {x} _ {\mathrm {t r g}} \| _ {2} ^ {2} ].
$$
We also utilize classifier-free guidance [21] to improve sample quality. We jointly learn an unconditional network,  $\hat{\mathbf{x}}_{\theta}(\mathbf{z}_t,t)$ , and a conditional network,  $\hat{\mathbf{x}}_{\theta}(\mathbf{z}_t,t,E(\mathbf{w}_{\mathrm{src}}))$ , by dropping the conditioning information with probability  $p = 0.1$  during training. When we drop the conditioning information, we cross-attend to a learnable embedding instead of the embedded source text. During sampling, we use guidance weight  $w$  and compute the prediction as
$$
\tilde {\mathbf {x}} _ {t} = w \hat {\mathbf {x}} _ {\theta} (\mathbf {z} _ {t}, t, E (\mathbf {w} _ {\mathrm {s r c}})) + (1 - w) \hat {\mathbf {x}} _ {\theta} (\mathbf {z} _ {t}, t).
$$
Setting  $w = 1.0$  corresponds to the conditional diffusion model while setting  $w > 1.0$  strengthens the influence of the conditioning information. We use  $w = 2.0$  for the seq2seq tasks and ablate this choice in section 5.
We can also generate multiple outputs  $\mathcal{S}$  for each input by sampling different latents  $\mathbf{z}_1\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ . We then select the most promising candidate with Minimum Bayes Risk (MBR) Decoding [15, 34]. In MBR decoding, we define a loss function  $\mathcal{L}$ , such as the negative Rouge, and use it to select a candidate  $\mathbf{w}_{\mathrm{MBR}} = \mathrm{argmin}_{\mathbf{w}\in \mathcal{S}}\frac{1}{|\mathcal{S}|}\sum_{\mathbf{w}'\in \mathcal{S}}\mathcal{L}(\mathbf{w},\mathbf{w}')$ . In our experiments, we use  $|\mathcal{S}| = 5$  and denote the results from using MBR decoding as MBR-5. We also report results using the ground truth to select the best candidate  $\mathbf{w}_{\mathrm{oracle}} = \mathrm{argmin}_{\mathbf{w}\in \mathcal{S}}\mathcal{L}(\mathbf{w},\mathbf{w}_{\mathrm{trg}})$  to provide an upper bound on the performance of our method given optimal sample selection. Because this requires knowledge of the ground-truth target text, we refer to this as Oracle sampling.
4 Datasets
We evaluate LD4LG on a variety of natural language datasets. ROCStories [42] is a corpus of 98k five-sentence commonsense stories, that capture casual and temporal relations. The AG News Topic
Classification [60] dataset consists of news articles across four topics: World, Sports, Business, Sci/Tech with article titles and descriptions from 120k training instances. We focus on generating the article descriptions in this work. The XSum [44] dataset consists of BBC articles from 2010 to 2017 covering a wide range of topics (e.g., News, Politics, Sports, etc.). The training split has 204k instances and each example contains a document and a summary. The QQP [9] dataset consists of 400k question pairs, where each example is two similar questions and a binary value indicating whether the two questions have the same meaning. The WMT 2014 English-German [5] dataset is a widely used machine translation dataset consisting of roughly 4.5 million sentence pairs. We present detailed dataset statistics in the appendix.
4.1 Evaluation Metrics.
We use MAUVE Score [47] and Perplexity (Ppl) to evaluate the quality of our generated text. MAUVE Score is a metric for open-ended text generation that compares the distribution of generated text with that of reference text using divergence frontiers. We follow Pillutla et al. [47] and use the GPT-2-Large model [49] to embed the text. Perplexity measures how likely the generated samples are according to an autoregressive language model; we use GPT-2-Large to compute perplexity.