SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yizhe Zhang, Jiatao Gu, Zhuofeng Wu, Shuangfei Zhai, Josh Susskind, Navdeep Jaitly Apple
{yizzhang, jgu32, zhuofeng_wu, szhai, jsusskind, njaitly}@apple.com
Autoregressive models for text sometimes generate repetitive and low-quality output because errors accumulate during the steps of generation. This issue is often attributed to exposure bias – the difference between how a model is trained, and how it is used during inference. Denoising diffusion models provide an alternative approach in which a model can revisit and revise its output. However, they can be computationally expensive and prior efforts on text have led to models that produce less fluent output compared to autoregressive models, especially for longer text and paragraphs. In this paper, we propose PLANNER, a model that combines latent semantic diffusion with autoregressive generation, to generate fluent text while exercising global control over paragraphs. The model achieves this by combining an autoregressive “decoding” module with a “planning” module that uses latent diffusion to generate semantic paragraph embeddings in a coarse-to-fine manner. The proposed method is evaluated on various conditional generation tasks, and results on semantic generation, text completion and summarization show its effectiveness in generating high-quality long-form text in an efficient manner.
Prefix: Horrible hotel. Horrible hotel. Seriously, horrible hotel.	1-gram	2-gram	3-gram
FT (Greedy decoding): The hotel was a nightmare. The rooms were filthy, the bathrooms were filthy, the food was disgusting, the food was disgusting. The food was disgusting. The food was disgusting. The food was disgusting....	The: 100.00%	The hotel: 100.00%	The hotel was: 100.00%
FT (Top-p sampling): Terrible, awful hotel. Horrible, awful hotel. I'm going to the awful hotel. I'll never get in. What the hell is this hotel? It's awful. I don't like it. Where? It's awful. I think they have a bathroom there. I'm going to the awful hotel. I'll never get in. I want out of here. This hotel is horrible. I'm going to the awful hotel. I'll never get in. I want out of here...	I: 13.39%
Horrible: 12.20%
The: 8.27%
Terrible: 5.91%
This: 4.33%	Horrible hotel: 12.50%
Terrible hotel: 6.05%
This is: 2.82%
I hate: 1.61%
I mean: 1.61%	Horrible hotel. Horrible: 3.48%
Horrible hotel. I: 2.61%
Terrible hotel. Terrible: 2.17%
Terrible hotel. Horrible: 1.30%
I could not: 0.87%
Ours (Greedy decoding): Staff was unhelpful. So rude from the front desk. The hotel tried to lose the reservation of my family which I had with the hotel. This hotel is way overpriced. Hallways smelled bad. Carpet was disgusting. I'll never go back. I don't recommend the hotel. I can stay anywhere else!!	I: 10.55%
The: 5.08%
We: 3.52%
Hotel: 2.73%
They: 2.73%	I was: 1.56%
Had to: 1.56%
The room: 1.17%
Dirty and: 1.17%
I called: 1.17%	The room was: 1.18%
The walls were: 0.78%
I was told: 0.78%
I called the: 0.78%
I booked through: 0.78%
Figure 1: Left: With a repetitive prompt, the finetuned GPT-2 large model (774M, FT) is still attracted to self-reinforced repetition (highlighted text) even under top-p sampling  $(\mathrm{K} = 50, \mathrm{p} = 0.92)$ . Right: the most frequent first n-grams of the generations for each method, derived from 512 generation roll-outs. Our proposed method results in a more diversified generation robust to the ill-composed prompt.
Autoregressive models trained with a teacher forcing strategy (Williams & Zipser, 1989) are considered the gold standard for text generation. However, a significant drawback of this approach is that it lacks the ability to correct any mistakes made during the generation process which can lead to errors that accumulate as the generation progresses. Previous work (Ott et al., 2018; Holtzman et al., 2019; Welleck et al., 2019; Xu et al., 2022) has observed that deterministic decoding methods have a
tendency to generate consecutive repetitions at the word, phrase and sentence levels. For example, with repetitive prompt, the model can enter an absorbing state where it produces repetitive outputs with higher and higher confidence by self-reinforcing the pattern (Xu et al., 2022) (Fig. 1). Through our experiments, we have observed that such degeneration is more prevalent in open-ended tasks that allow the model greater freedom for creativity. Even for large language models, the generation can drift away from the desired semantics, especially when the model is poorly prompted or has high initial probabilities (Xu et al., 2022).
Why does using maximum likelihood decoding lead to repetitions during generation which is significantly different from the training data distribution? One possible explanation for this is "exposure bias" (Bengio et al., 2015) arising from the discrepancy between the training and inference phases in the teacher forcing training strategy. During training phase, the model focuses only on predicting the next token. However, during inference, predicting the next token alone can be myopic because the model may not have enough foresight to anticipate its impact on future generation. This can also be seen as the "distribution shift" issue of behavior cloning (de Haan et al., 2019), where the model is trained to mimic the expert's actions on the states encountered by the expert in the training data. However, small differences between the model and the expert can compound over multiple steps, leading the model to states it never encountered during training, rendering unreliable and undesirable predictions.
Although many approaches have been proposed to address this issue, such as adversarial models (Yu et al., 2017; Lamb et al., 2016; Zhang et al., 2017), reinforcement learning (Li et al., 2016b) or repetition penalties (Xu et al., 2022), they attempt to improve the global aspects of the generation by making local adjustments which still follow the autoregressive generation recipe. Diffusion models provide an alternative solution – the model can revisit and revise its output iteratively, potentially rendering more global control of the generation in a non-autoregressive manner. However, these text diffusion models can generate less fluent text compared to autoregressive ones (Gong et al., 2023). Also, when generating long text, the diffusion process involves multiple passes of the underlying denoising model over a long generation length, making it computationally expensive. The discrete nature of text also presents a challenge for diffusion models, which can suffer from “rounding errors” when converting between the text token and its embedding (Li et al., 2022; Lin et al., 2022).
Instead of performing diffusion on the original text or the corresponding word embeddings, we propose to apply diffusion techniques to the latent semantic space (Rombach et al., 2022; Lovelace et al., 2022). To achieve this, we learn a fixed number of continuous semantic tokens that encode salient information at the paragraph level. These tokens can then be used to reconstruct the original text. The latent diffusion can be additionally conditioned on an external signal to generate the semantic tokens. Finally, a decoder maps the obtained semantic tokens back to the raw text space. This process combines a non-autoregressive semantic diffusion approach with an autoregressive decoding technique. The semantic diffusion process handles the "planning", enabling the modification of semantics in a coarse-to-fine manner, while the decoder handles the "decoding" by translating the semantics into raw text, with less flexibility in controlling the meaning. We call our proposed method PLANNER (Paragraph-levelL Diffusion model for Embedding Representation).
Our contributions include: (i) We propose a latent semantic diffusion model for paragraphs that incorporates both non-autoregressive semantic diffusion and autoregressive generation. This allows us to generate fluent text while being able to exercise global control inherited from a diffusion model. (ii) We study the essential requirements for a good latent space for paragraph diffusion models. (iii) We evaluate the effectiveness of our proposed method on various conditional generation tasks. Thanks to the iterative refinement of desnoising diffusion, our method enjoys less repetitive and more diverse generation, while maintaining good fluency and relevancy, comparing with autoregressive and text diffusion baselines (Li et al., 2022; Lin et al., 2022).
Diffusion Probabilistic Models The standard diffusion model (DM) (Ho et al., 2020; Song & Ermon, 2019) learns the data distribution  $p(x)$  by gradually denoising a normally distributed variable in a Markov chain of length  $T$ . The diffusion process can be viewed as a continuous-time stochastic process (Song et al., 2021b; Kingma et al., 2021) where the initial data point  $x \in \mathbb{R}^N$  is progressively corrupted by noise according to a predefined signal-noise schedule  $\{\alpha_t, \sigma_t\}$ , resulting in time-
dependent corrupted data  $\{\pmb{x}_t|t\in [0,1],\pmb{x}_0 = \pmb{x}\}$ . The transition distribution is given by:
$$
q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {s}\right) = \mathcal {N} \left(\boldsymbol {x} _ {t}; \alpha_ {t | s} \boldsymbol {x} _ {s}, \sigma_ {t | s} ^ {2} I\right), \tag {1}
$$
where  $\alpha_{t|s} = \alpha_t / \alpha_s$ ,  $\sigma_{t|s}^2 = \sigma_t^2 - \alpha_{t|s}^2\sigma_s^2$ , and  $s < t$ . When  $\pmb{x}_s = \pmb{x}$ , the marginal distribution  $q(\pmb{x}_t|\pmb{x})$  is given as  $q(\pmb{x}_t|\pmb{x}) = \mathcal{N}(\pmb{x}_t; \alpha_t\pmb{x}, \sigma_t^2 I)$ . The diffusion model relies on a parametric function  $\theta$  optimized to reverse the diffusion process by denoising  $\pmb{x}_t$  to the clean input  $\pmb{x}$ . The model is trained using a weighted reconstruction loss:
$$
\mathcal {L} (\theta) = \mathbb {E} _ {\boldsymbol {x} _ {t} \sim q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x}\right), t \sim [ 0, 1 ]} \left[ \omega_ {t} \cdot \| \boldsymbol {F} _ {\theta} \left(\boldsymbol {x} _ {t}, t\right) - \boldsymbol {x} \| _ {2} ^ {2} \right], \tag {2}
$$
where  $\omega_{t} = \alpha_{t}^{2} / \sigma_{t}^{2}$ , (s.t.  $\alpha_{t}^{2} + \sigma_{t}^{2} = 1$ ) is the signal-to-noise-ratio (SNR) and  $F_{\theta}(\cdot)$  denotes the backbone denoising function. Sampling from the learned model can be performed using either ancestral sampling (DDPM) (Ho et al., 2020) or a deterministic DDIM sampler (Song et al., 2021a). While the DM is capable of generating high-quality samples, the fact that the corrupted data  $\boldsymbol{x}_{t}$  shares the same space as the input  $\boldsymbol{x}$  results in inefficient training (Jing et al., 2022) and difficulty in learning abstract and semantically meaningful latent spaces (Preechakul et al., 2022).
Latent Diffusion Models To improve the efficiency, the Latent Diffusion Model (LDM) (Rombach et al. 2022) introduces an explicit separation between the compressive and generative learning phases of training diffusion models. It employs an autoencoding model consisting of an encoder  $\mathcal{E}(\cdot)$  and a decoder  $\mathcal{D}(\cdot)$  to learn a low-dimensional latent space that is perceptually equivalent to the image space when decoded, but with reduced computational complexity, while retaining the perceptual quality of generated samples. The reweighted objective for training LDM is given by:
$$
\mathcal {L} (\theta) = \mathbb {E} _ {\boldsymbol {z} _ {t} \sim q \left(\boldsymbol {z} _ {t} \mid \boldsymbol {z}\right), \boldsymbol {z} = \mathcal {E} (\boldsymbol {x}), t \sim [ 0, 1 ]} \left[ \omega_ {t} \cdot \| \boldsymbol {F} _ {\theta} \left(\boldsymbol {z} _ {t}, t\right) - \boldsymbol {z} \| _ {2} ^ {2} \right], \tag {3}
$$
where  $z$  is obtained from  $\mathcal{E}$  during training. The generated  $z$  can be decoded to image using  $\mathcal{D}$ .
Text diffusion models Early attempts on using diffusion models for discrete data used a noising processes which masked or randomly mutated the discrete tokens (Austin et al., 2021; Hoogeboom et al., 2021). Recently, Diff-LM (Li et al., 2022) and DiffuSeq (Gong et al., 2023) have instead used a continuous token embedding space, converting the continuous token embeddings to text via "rounding". Analog Bits (Chen et al., 2022) converts raw text into a set of bits and models them as analog bits with a continuous diffusion model. (Lovelace et al., 2022) performed diffusion model on the contextualized BART embeddings rather than on the word embedding space. (Zhu et al., 2022) has applied text diffusion to image captioning and achieved good performance.
However, existing text diffusion models present several issues: (i) The varying length of the input text necessitates the prediction of additional length or superfluous paddings, and (ii) token generation in parallel may result in disfluent text and/or frequent repetitions especially when the generation is long. We instead employ the diffusion model to learn paragraph embeddings that contain fewer fixed-sized tokens, which allows for computational benefits and improved fluency.
Text Variational Autoencoders Text VAEs (Bowman et al., 2016; Kim et al., 2018; Li et al., 2020) have been particularly useful for learning a smooth and interpretable representation space, as well as for generating diverse text. However, one of the challenges is the KL vanishing problem (Bowman et al., 2016), which results in the decoder disregarding the latent code sampled from the prior distribution during the inference stage. Our approach can be perceived as to address this issue by leveraging a more flexible prior distribution to ensure the codes can strongly influence the output text distribution.
4 PLANNER: A Language Diffusion Model on Paragraph Embeddings
We use latent diffusion to improve the diversity and fluency of paragraphs generated from the model. Our model comprises two parts (Fig. 2) - a paragraph embedder via variational autoencoder (VAE) that learns a meaningful and smooth latent space that corresponds to the original text space, and a diffusion model that generates latent codes corresponding to the semantics of longer text.
Figure 2: Model overview. Left: a variational paragraph embedder is learned to encode paragraph into a fixed amount of latent codes. Right: the latent diffusion model based on transformer block is applied to generate the latent codes. The decoder finally translates them into the text. (BOS: Begin of Sentence token, EOS: End of Sentence token)
4.1 Learning a Variational Paragraph Embedder
Instead of applying diffusion to tokens directly to generate long text, we propose to learn a set of latent codes  $\mathbf{z} = \{z^{(1)},\dots ,z^{(k)}\} \in \mathbb{R}^{k\times h}$ , which we call paragraph embeddings, that capture the semantics in the target text (of length up to 512 tokens), where  $h$  denotes the embedding dimension. These paragraph embeddings have shorter length, such as  $k = 16$ , than the original text.
To obtain such embeddings  $z$ , we train a transformer-based encoder-decoder model. The architecture used for the autoencoder is shown in Fig. 2. The encoder  $\mathcal{E}$  and decoder  $\mathcal{D}$  construct a bidirectional mapping between the discrete data space and the latent code space. The paragraph embeddings  $z$  are extracted by taking the first  $k$  hidden state vectors of dimension  $h$  from the final layer of  $\mathcal{E}$ , which are fed into the initial steps of the decoder which is trained to reconstruct the original text. It's worth noting that the paragraph embeddings share the same hidden dimension  $h$  as the word embeddings, and forming a manifold in the word embedding space. Pretrained BERT and GPT-2 models are used to initialize  $\mathcal{E}$  and  $\mathcal{D}$ , respectively. The manifold of the learned embeddings ideally possesses several desirable properties, including low conversion error, local smoothness and distributional smoothness.
Conversion error Ideally, the original input  $\pmb{x}$  can be perfectly reconstructed via  $\hat{\pmb{x}} = \mathcal{D}(\pmb{z})$ ,  $\pmb{z} = \mathcal{E}(\pmb{x})$ , and modeling the lower-dimensional continuous space  $p(z)$  is equivalent to modeling  $p(x)$ . However, in practice a loss of information can occur when converting raw text into paragraph embeddings or when doing the reverse. We assess the conversion loss by computing the BLEU score  $(\mathrm{BLEU}_{\mathrm{clean}})$  between the input  $\pmb{x}$  and the reconstruction  $\hat{\pmb{x}}$ .
Local smoothness To generate target text that is fluent and consistent with the corresponding paragraph embeddings, it is essential to achieve a certain level of local smoothness in the paragraph embeddings space. Ideally, a slight variation in the input vector  $\pmb{x}$  would not cause a significant change in the resulting encoded vector  $\mathcal{E}(\pmb{x})$ . Similarly, a small perturbation in the latent vector  $\pmb{z}$  should not lead to a significant change in the decoded vector  $\mathcal{D}(z)$ . Otherwise, the error accumulated in the diffusion process when generating  $z$  could result in an inaccurate realization of the desired semantics. To accomplish this, the denoising autoencoder is trained by substituting (Sub) input tokens with random tokens with probability  $p$ . The local smoothness is measured using the BLEU score  $(\mathbf{BLEU}_{\mathrm{robust}})$  between the input  $x$  and the denoising output from corrupted input  $\mathcal{D}(\mathcal{E}(\tilde{\pmb{x}}))$ , where  $\tilde{\pmb{x}} = \mathrm{Sub}(\pmb{x}, p = 0.3)$ . The level of injected noise will affect both the conversion error and the local smoothness, and it is important to strike a balance between the two.
Distributional Smoothness The diffusion model may face difficulty in learning a distribution,  $p(z)$  that is highly multimodal, or has density that are associated with a large Lipchitz constant (i.e., has abrupt changes). Therefore, we employ a text VAE (Bowman et al. 2016; Li et al. 2020) to encourage the posterior to take on a form close to a Gaussian distribution. Specifically, we parameterize  $q(z|x)$  to be  $\mathcal{N}(\mathcal{E}_{\mu},\mathcal{E}_{\nu})$  and maximize the objective  $\mathcal{L}(\mathcal{E},\mathcal{D};\boldsymbol {x}) = \mathbb{E}_{q(\boldsymbol {z}|\boldsymbol {x})}[\log p(\boldsymbol {x}|\boldsymbol {z})] - \beta \cdot \mathrm{KL}(q(\boldsymbol {z}|\boldsymbol {x})||p(\boldsymbol {z}))$  Here  $\mathcal{E}_{\mu}$  and  $\mathcal{E}_{\nu}$  represent the posterior mean and variance predictions of the encoder  $\mathcal{E}$ , while the
hyperparameter  $\beta$  controls the strength of regularization. It is typically set to a small value to alleviate the notorious posterior collapsing issue (Bowman et al., 2016) in text VAE. To gauge the distributional smoothness of the paragraph embedding space, we select two examples,  $x$  and  $x^{\prime}$  at random from the training set and interpolate their embeddings to compute  $z_{\mathrm{INT}} = \frac{1}{2}\mathcal{E}(\boldsymbol {x}) + \frac{1}{2}\mathcal{E}(\boldsymbol{x}^{\prime})$ . We then evaluate the perplexity  $(\mathbf{PPL}_{\mathrm{int}})$  of the decoded interpolation  $\mathcal{D}(z_{\mathrm{INT}})$  using a GPT-2 model.
4.2 Planning then Decoding: A Latent Diffusion Model for Paragraph Embeddings
Training phase We now use the learned mean paragraph embeddings  $\mathbf{z} = \mathcal{E}_{\mu}(\mathbf{x})$  to train a continuous-time latent diffusion model as in Fig. 2 while keeping  $\mathcal{E}$  and  $\mathcal{D}$  frozen. We conducted experiments using two types of conditioning signal: (i) class labels, such as positive or negative sentiment labels, and (ii) raw text, such as preceding context or the document to be summarized. For class labels, we learned a label embedding  $\mathbf{y} \in \mathbb{R}^h$  to represent each class. For the raw text, we applied a conditional feature encoder  $\tau$  to the input and used the hidden states at the last layer as  $\mathbf{y} \in \mathbb{R}^{c \times h}$ , where  $c$  represents the number of feature embeddings.
During training, we gradually add noise to  $z$  via a cosine scheduler (Ho et al., 2020), and use a signal prediction scheme as the training objective (Kingma et al., 2021). For our denoising backbone model  $F(\cdot)$ , we use a transformer block similar to the one in the DiT (Peebles & Xie, 2022) model. Specifically, we fed  $y$  and the time embedding  $t \in \mathbb{R}^h$  into the model through two channels, namely cross-attention and adaptive layer norm (adaLN) (Peebles & Xie, 2022). For the cross-attention, the conditional embeddings  $t$  and  $y$  are concatenated into a sequence of length  $c + 1$ . The transformer block is modified to enable multi-head cross-attention to the conditional embeddings.
For the adaLN, we flattened and projected  $\mathbf{y}$  to a vector of  $\mathbb{R}^h$  using a linear projection layer. We then added the projected  $\mathbf{y}$  to the time embedding  $t$ . Instead of directly learning dimension-wise scale and shift parameters  $(\beta$  and  $\gamma)$  in the standard Layer Norm (LN), these parameters are regressed from the sum of the embeddings. In addition, dimension-wise scaling parameters  $\alpha$  are regressed and applied immediately before any residual connections within the transformer block. This has been shown to be efficient and effective in image diffusion models (Peebles & Xie, 2022).
Inference phase During the inference process, we start with random Gaussian embeddings and use a fixed number of steps  $T$  to generate the final  $\pmb{z}$ . The resulting embeddings are then used as inputs for  $\mathcal{D}$  to generate the text using a deterministic decoding method like greedy decoding. We provide discussion and ablation study on using stochastic decoding in the App.
We applied the classifier-free guidance (CFG) (Ho & Salimans, 2021) during the inference steps. After each step, we apply a dynamic thresholding technique that was introduced in Imagen (Saharia et al., 2022) for post-processing. However, we do not use the rescaling step in Imagen because rescaling step can completely alter the underlying semantics, as we have not imposed any constraints to ensure that the generated output remains the same after rescaling the latent code. In contrast, for Imagen, where the generation takes place in the raw pixel space, rescaling will predominantly retain the shape information while altering only the contrast and brightness.
5 Experimental Setups
We tested the effectiveness of our model in three different conditional generation tasks including sentiment-guided generation, long-form text completion, and summarization. The tasks can require generating text of hundreds of tokens in length, making them suitable to assess model performance.
Datasets For the Sentiment-guided generation task, we used the TripAdvisor dataset provided by (Li et al., 2014). By exclusively selecting reviews with a rating of 1 or 5 and balancing the two ratings via subsampling, we acquired 218,509 reviews. For the text completion task, our model was assessed on two datasets: 1) the aforementioned TripAdvisor review dataset with postprocessing to remove reviews that are less than 20 or more than 512 tokens, result in 690,862 samples, and 2) one-tenth of the overall C4 datasets (Raffel et al., 2020), which contains 36.5M samples. For each sample, we extracted the initial two sentences from a paragraph as source context, and predicted the remainder of the text as target. The datasets were partitioned into training, validation, and test in the ratios of
The aim of  $\mathcal{D}$  is to accurately convert the  $z$  into a meaningful text, thus deterministic decoding is desirable.
(0.96, 0.02, 0.02). For the summarization task, we use CNN/DailyMail (Hermann et al., 2015) and XSum (Narayan et al., 2018).