SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\boldsymbol {\mu} _ {\theta} (\mathbf {z} _ {t}, s, t) = \boldsymbol {\mu} _ {Q} (\mathbf {z} _ {t}, \hat {\mathbf {x}} _ {\theta} (\mathbf {z} _ {t}, t), s, t), \quad \sigma^ {2} (s, t) = 1 - \alpha_ {t | s}.
$$
We therefore substitute our estimate of the clean data into the posterior distribution of  $q(\mathbf{z}_s|\mathbf{z}_t,\mathbf{x})$  to parameterize the mean of our generative process  $p_{\theta}(\mathbf{z}_s|\mathbf{z}_t)$ . We follow Ho et al. [22] and set the variance of  $p_{\theta}(\mathbf{z}_s|\mathbf{z}_t)$  to  $\sigma^2 (s,t) = 1 - \alpha_{t|s}$ , a choice given by the variance of the forward process.
For generation, we utilize the standard DDPM sampler, also known as the ancestral sampler [22]. We sample some initial noise  $\mathbf{z}_{t_1} = \mathbf{z}_1\sim \mathcal{N}(\mathbf{0},\mathbf{I})$  and iteratively apply the update rule
$$
\mathbf {z} _ {t _ {i + 1}} = \boldsymbol {\mu} _ {\theta} \left(\mathbf {z} _ {t _ {i}}, t _ {i + 1}, t _ {i}\right) + \sigma \left(t _ {i + 1}, t _ {i}\right) \epsilon
$$
where  $\epsilon \sim \mathcal{N}(\mathbf{0},\mathbf{I})$  and the intermediate timesteps  $1 = t_{1} > t_{2} > \ldots >t_{T} = 0$  linearly interpolate between 1 and 0. We use  $T = 250$  sampling timesteps by default.
B Additional Language Autoencoder Results
We present results for our language autoencoders on XSum, QQP and WMT14 in Table 9. We observe that our proposed language autoencoders are similarly effective for these datasets.
We also ablate the performance as we vary the dimensionality of the latent space in Table 10. We observe, as expected, that the reconstruction performance improves as the dimensionality of the latent space increases and degrades as we decrease the size of the latent representation. We found our default dimensionality of  $32 \times 64$  to be generally effective for high quality reconstructions across datasets.
Table 9: Effectiveness of Language Autoencoder
Method	Latent Dimensions	Hidden Units	XSum	QQP
Rouge-1/2/L	BLEU	Rouge-1/2/L	BLEU
BART-Base	L × 768	≤ 49,152	99.9/99.9/99.9	99.9	99.9/99.9/99.9	99.8
BART-Base Autoencoder	32 × 64	2048	99.8/99.6/99.8	99.3	99.9/99.8/99.9	99.1
FLAN-T5-Base	L × 768	≤ 49,152	65.7/51.2/59.9	45.7	26.9/14.1/24.3	10.8
FLAN-T5-Base Autoencoder	32 × 64	2048	99.6/99.3/99.6	98.8	99.8/99.5/99.7	98.5
WMT14 English	WMT14 German
Rouge-1/2/L	BLEU	Rouge-1/2/L	BLEU
MT5-Base Autoencoder	32 × 64	2048	99.7/99.2/99.7	99.2	99.8/99.4/99.8	99.1
Table 10: Ablation of Autoencoder Latent Dimensionality
Method	Latent Dimensions	Hidden Units	RocStories
Rouge-L	BLEU
BART-Base	L × 768	≤ 49,152	98.8	97.5
BART-Base	32 × 32	1024	97.0	92.4
Autoencoder	32 × 64	2048	99.2	97.6
64 × 64	4096	99.2	97.7
C Impact of Sampling Steps
We present the results from different sampling configurations for the ROCStories dataset in Table 11. We also report the wall clock time needed to generate the 1000 samples across the different numbers of sampling timesteps while batching the generations with a batch size of 128.
We find that the number of sampling steps introduces a tradeoff between the diversity and the quality of the text, with more sampling steps leading to more fluent but less diverse text and fewer sampling steps leading to less fluent but more diverse text. When using BART-base, the MAUVE score is maximized when utilizing only 100-250 steps, demonstrating that it achieves a reasonable balance between diversity and quality. When utilizing FLAN-T5-base, on the other hand, we find that the MAUVE score improves monotonically with increased sampling steps. This suggests that the latent distribution of the FLAN-T5-base autoencoder may be more challenging to learn. Increasing the capacity of the denoising network or the language autoencoder may therefore be beneficial when using FLAN-T5-base.
We observe that the sampling time scales with the number of sampling steps as expected, although there is also a fixed cost from the reconstruction network and the autoregressive decoder that is independent of the number of sampling steps.
Table 11: Evaluation of different sampling configurations. We use 250 steps by default.
ROCStories
Sampling Steps	MAUVE ↑	Ppl ↓	Div ↑	Mem ↓	Wall Clock Time (1000 samples)
Reference	-	.951.007	21.1 .3	.414.003	.362.003	-
LD4LG (BART-base)	50	.684.031	52.6 .3	.407.004	.337.003	1m27s
100	.719.022	38.5 .8	.368.002	.392.001	1m55s
250	.716.019	30.6 .5	.331.005	.441.004	3m20s
500	.704.033	28.1 .3	.313.003	.462.003	5m44s
1000	.667.026	25.9 .1	.295.002	.481.004	10m30s
LD4LG (FLAN-T5-base)	50	.331.028	67.9 .7	.456.001	.283.001	1m34s
100	.421.012	48.7 .7	.423.002	.334.002	2m02s
250	.481.007	37.5 .4	.389.002	.387.002	3m29s
500	.495.024	32.8 .6	.370.006	.413.006	5m51s
1000	.522.023	30.6 .3	.360.004	.432.005	10m38s
Table 12: Evaluation of different nucleus sampling configurations.
ROCStories
Sampling Parameter (p)	MAUVE ↑	Ppl ↓	Div ↑	Mem ↓
GPT-2-Medium	.90	.762 .027	19.6 .3	.362 .008	.718 .006
.95	.788 .025	20.0 .2	.372 .002	.688 .006
.98	.782 .020	20.2 .3	.378 .002	.666 .008
1.00	.793 .024	20.5 .4	.385 .004	.637 .006
D GPT-2 Sampling Ablation
We report an ablation of the nucleus sampling parameter,  $p$ , in Table 12. The memorization does exhibit some sensitivity to the nucleus sampling parameter, but the memorization is consistently higher than the LD4LG models across all sampling configurations.
All of the models presented in this work are trained on a single Nvidia A6000 except for the DiffuSeq XSum baseline which was trained with two Nvidia A6000s.
E.1 Language Autoencoders
We adopt the pre-LN design [70] for both the compression and reconstruction networks and therefore apply layer normalization before all attention and feedforward blocks. We also adopt query-key normalization [11] and apply RMSNorm [74] to the queries and keys before computing the dot product similarities in the attention mechanism. We found that this enabled training with a larger learning rate which accelerated training. We present the hyperparameters for our language autoencoders across all datasets in this work in Table 13.
We also report additional details such as the number of trainable parameters and training time. The training time is similar across datasets because we use the same hyperparameters, so we simply report the training times for the ROCStories dataset for the monolingual models. For the MT5-base base autoencoder, we report the training time for the German autoencoder which is similar to the English autoencoder. We note that our implementation is not optimized for runtime and that pre-computing and caching the language encoder representations would significantly accelerate training.
E.2 Latent Diffusion For Language Generation
We present the training details across the different datasets in Table 14. We tuned hyperparameters using the validation MAUVE scores for the ROCStories dataset and found that they generally transferred well across datasets. We therefore used the same hyperparameters across datasets, except that we utilized the L1 loss instead of the L2 loss for the Seq2Seq tasks. Consistent with prior work on image-to-image diffusion models [53], we observed that the L1 loss improved the fidelity of the generations at the cost of sacrificing some diversity. This improved fidelity translated to improvements in our metrics of interest, although the L2 loss may still be desirable for settings where diversity is of greater importance. For the unconditional and class-conditional language models, we did not observe overfitting to be a problem and simply use the final checkpoint for evaluation. For the monolingual Seq2Seq tasks, we utilize the checkpoint with the best validation ROUGE-L. For machine translation, we utilize the checkpoint with the best validation SacreBLEU.
For the machine translation experiments, we observed benefits from rescaling the noise schedule to emphasize training at higher levels of noise. This idea was introduced by Hoogeboom et al. [27] and Chen [7] to improve high-resolution image diffusion models. Both Hoogeboom et al. [27] and Chen [7] shift an existing noise schedule by some scale factor,  $s$ , to increase the time spent at higher noise levels. Given a noise schedule  $\alpha_{t}$  with SNR  $\lambda_{t} = \frac{\alpha_{t}^{2}}{1 - \alpha_{t}^{2}}$ , the shifted noise schedule,  $\alpha_{t,s} \in [0,1]$ , is
Table 13: Training details for our language autoencoders.
Language Model
BART-base	FLAN-T5-base	MT5-Base
Trainable Params	26M	26M	591M
Compression Architecture	Perceiver Resampler [2]
Perceiver Layers	3	3	1
Perceiver Dimension	768	768	768
Self-Attention Heads	12	12	12
Autoencoder Latent Length (l)	32	32	32
Autoencoder Dimension (dae)	64	64	64
Reconstruction Architecture	Transformer [69]
Transformer Layers	3	3	1
Transformer Dimension	768	768	768
Self-Attention Heads	12	12	12
Activation Function	GELU [20]
Max Seq Length	64	64	128
Optimizer	AdamW [39]
Learning Rate	1e-4	1e-4	1e-4
(β1, β2)	(0.9, 0.999)	(0.9, 0.999)	(0.9, 0.999)
Batch Size	256	256	128
Warmup Steps	1000	1000	1000
Learning Rate Schedule	Linear Decay
Weight Decay	1e-2	1e-2	1e-2
Gradient Clipping	1.0	1.0	1.0
Training Steps	50k	50k	50k
Training Time	12h38m	20h17m	20h29m
defined so that
$$
\frac {\alpha_ {t , s} ^ {2}}{1 - \alpha_ {t , s} ^ {2}} = \lambda_ {t, s} = \lambda_ {t} * s ^ {2} = \frac {\alpha_ {t} ^ {2}}{1 - \alpha_ {t} ^ {2}} * s ^ {2}.
$$
Given  $\alpha_{t}$  and the scale factor  $s$ , the scaled noise schedule  $\alpha_{t,s}$  can be computed in closed-form. Using the relationship that  $\alpha_{t}^{2} = \mathrm{sigmoid}(\log (\lambda_{t}))$  (see Kingma et al. [30]), the new noise schedule can be computed as
$$
\alpha_ {t, s} ^ {2} = \operatorname {s i g m o i d} \left(\log \left(\lambda_ {t, s}\right)\right) = \operatorname {s i g m o i d} \left(\log \left(\lambda_ {t} * s ^ {2}\right)\right) = \operatorname {s i g m o i d} \left(\log \left(\lambda_ {t}\right) + 2 \log (s)\right).
$$
We employ a shifted cosine noise schedule with  $s = 0.1$  for machine translation. Past work on text diffusion for machine translation observed that training at higher levels of noise improves the models utilization of the conditioning information (i.e. the source sentence) [73].
During the inference process, image diffusion models typically rescale the estimate of the data to the range of pixel values (i.e. [-1,1]) at each sampling step. When we restrict the latent space so that  $\|\mathbf{x}_i\|_2^2 = d_{\mathrm{ae}}$ , we similarly rescale the intermediate estimates of the data to enforce this constraint. This design decision is not critical and similar performance is achieved without this rescaling. We did, however, observe that this made the generative process more robust to large guidance weights which may be important in some settings. This observation is consistent with prior findings from text-to-image diffusion [55].
We also report the wall clock times for training the models, although our implementation could be further optimized to improve training times. The primary cause of the slowdown for AG News compared to ROCStories, for instance, stems from additional validation sampling and logging for class-conditional generation during training.
When decoding the sampled latent vectors, we utilize beam search with a beam size of 4, a repetition penalty of 1.2 [29], and prevent generations of duplicate trigrams.
E.3 BART-Diffusion
For our BART-Diffusion baseline, we utilize the same denoising architecture as our LD4LG method. As discussed in the main paper, the sequence length of the BART features vary with the length of
the input text. During training, the sequence length is simply determined by the training instance. To select the length of the Gaussian noise during generation, we sample a length from the empirical distribution of lengths in the training set.
We observed that the v-prediction parameterization was less effective in this setting and the  $\epsilon$ -prediction parameterization was unstable. We therefore adopted the x-prediction parameterization. This is consistent with past work that has found the x-prediction parameterization to be more effective for high-dimensional data [36, 8].
Another challenge is that we can no longer control the scale of the latent space. We therefore follow common practices from latent image diffusion and normalize the latent space to have unit variance [51]. When normalizing the latent space, we utilize the first batch of training data to compute the mean for each feature dimension, averaging across the samples in the batch and the sequence lengths of the samples. Therefore, we compute the mean vector  $\hat{\pmb{\mu}} = \frac{1}{b\ell}\sum_{b,\ell}\mathbf{x}_{b,\ell},\hat{\pmb{\mu}}\in \mathbb{R}^{d}$  where  $\mathbf{x}\in \mathbb{R}^{b\times \ell \times d}$  is some batched data. We then compute the global variance across all dimensions in the centered latent space  $\hat{\sigma}^2 = \frac{1}{b\ell d}\sum_{b,\ell ,d}(\mathbf{x}_{b,\ell ,d} - \hat{\boldsymbol{\mu}}_d)^2,\hat{\sigma}^2\in \mathbb{R}$  to rescale the latent space to have unit variance. We otherwise train this baseline with the same hyperparameters as LD4LG.
E.4 Diffusion LM
We train our Diffusion-LM models utilizing the public implementation by Li et al.  $[36]^2$ . We utilize the provided command and hyperparameter settings for the ROCStories dataset. To adapt it to the AG News dataset, we increase the batch size from 64 to 128 and set the number of training steps to  $250k$  match our training configuration. We otherwise utilize the same hyperparameter settings as the ROCStories model. We attempted to double the learning rate from 1e-4 to 2e-4 to account for the doubled batch size, but observed training instabilities and therefore used the original learning rate of 1e-4.
E.5 GPT-2
We present the default hyperparameters for the GPT-2-Medium baseline in Table 15. For sampling from GPT-2, we prompt it with a B0S token and utilize nucleus sampling  $(p = 0.95)$ . We use the same repetition penalty of 1.2 [29] that we use for the LD4LG language decoders and similarly prevent generations of duplicate trigrams.
E.6 DiffuSeq
For the QQP dataset, we compute the metrics with the model generations released by Gong et al. [16]. We utilize the official implementation from Gong et al.  $[16]^3$  to train a DiffuSeq model on the XSum dataset. In their work, the DiffuSeq models were trained with the same hyperparameters across all datasets considered, except for the number of training steps which varied across datasets. We therefore adopt their default hyperparameters for the XSum dataset.
We observed that the DiffuSeq models were trained for much longer than our models. The official implementation also utilized gradient accumulation with microbatches of 128 to achieve a large effective batch size of  $4096^{4}$ . We trained the XSum DiffuSeq model for 960k iterations which is significantly longer than the 250k iterations used by our LD4LG XSum model. Due to the use of gradient accumulation, this corresponds to 30k gradient updates. The XSum DiffuSeq baseline was therefore trained for over  $3.8 \times$  more epochs than our method.
A limitation of the DiffuSeq model compared to LD4LG is that it concatenates the source and target sequences as the input to their transformer model. DiffuSeq therefore scales quadratically with respect to the combined length of the source and target sequence. Our denoising network, on the other hand, operates upon a fixed sequence length of  $\ell = 32$  latents and only cross-attends to the
$^{2}$ https://github.com/XiangLi1999/Diffusion-LM
<sup>3</sup>https://github.com/Shark-NLP/DiffuSeq
<sup>4</sup>We note that the original DiffuSeq implementation had a bug in its implementation of distributed training (see https://github.com/Shark-NLP/DiffuSeq/issues/37. We describe the behavior of the original implementation.
source representations. As a result, our method scales linearly with respect to the length of the source sequence<sup>5</sup>. This enables LD4LG to more efficiently incorporate long contexts than DiffuSeq.
By default, the official DiffuSeq implementation limits the combined length of the source and target sequences to a maximum length of 128. This could put it at a disadvantage compared to our model which incorporates up to 256 tokens of the source sequence. To ensure a fair comparison, we also experimented with increasing the maximum sequence length for the DiffuSeq model to 256 tokens, which significantly increases the training overhead. After training the model for 640k iterations, which took 5 days with two Nvidia A6000 GPUs, we observed worse performance than the model using the default length of 128.
E.7 Encoder-Decoder Language models
We report training hyperparameters for fine-tuning the pre-trained encoder-decoder language models models on the Seq2Seq datasets in Table 16. We perform early stopping with the validation ROUGE-L.
E.8 Evaluation Metrics
For the MAUVE, ROUGE, BLEU, BERTScore, Perplexity, and SacreBLEU metrics, we utilize the implementations provided by the Huggingface evaluate library (https://huggingface.co/docs/evaluate/). For SacreBLEU, we follow prior work and use the int1 tokenizer if the target language is German and use the 13a tokenizer if the target language is English.
For the n-gram metrics, we utilize the en_core_web_sm tokenizer from Spacy (https://spacy.io/) to split the generations into tokens.
ROCStories [42]. The dataset consists of 98,161 instances. We hold out 1,000 instances for validation, 4,000 instances for testing, and utilize the remaining 93,161 instances for training.
AG News Topic Classification [60]. The dataset consists of titles and short descriptions from news articles. We discard the titles and focus on generating the descriptions in this work. The official train/test splits have 120k training instances and 7,600 testing instances. We hold out 1,000 instances from the training set for validation. We therefore utilize 119k training instances, 1,000 validation instances, and 7,600 test instances.
XSUM [43]. The dataset consists of BBC articles from 2010 to 2017 covering a wide range of topics (e.g., News, Politics, Sports, etc.). Each example in the dataset consists of a news article and a summary. It has 204,045 training instances, 11,332 validation instances, and 11,334 test instances.
QQP [9]. The dataset consists of 400k question pairs, where example consists of two similar questions and a binary value indicating whether the two questions have the same meaning. The semantically similar questions can be utilized as a paraphrasing dataset. We use the version released by Gong et al. [16] to enable direct comparison. It has 144,715 training instances, 2,048 validation instances, and 2,500 test instances.
WMT 2014 English-German [5]. The dataset consists of roughly 4.5 million paired English and German sentences for training. The validation and testing splits each have roughly 3k paired sentences.
G Qualitative Examples
We present random unconditional samples from the diffusion models for the ROCStories (Table 17) and AG News (Table 18) datasets. We note that because the Diffusion-LM learns token representations from scratch and cannot model rare words, Li et al. [36] replace rare words with an UNK token. We observe that these tokens are often generated, leading to incoherent text. This problem is particularly
For LD4LG, the frozen language encoder still scales quadratically with the source sequence length, but the source representations can be pre-computed and cached prior to training.
Table 14: Training details for LD4LG across different datasets.
ROCStories	AG News	XSum	QQP	WMT14-En-De
Trainable Params	188M	190M	217M	217M	218M
Sampling Timesteps	250
Noise Schedule	Cosine	Cosine	Cosine	Cosine	Shifted Cosine (s=0.1) [27, 7]
Regression Loss	L2	L2	L1	L1	L1
Transformer Layers	12
Transformer Dimension	768
Self-Attention Heads	12
Dense Connections [4]	3
Activation Function	GeGLU [59]
Optimizer	AdamW [39]
Learning Rate	2e-4	2e-4	2e-4	2e-4	4e-4
(β1, β2)	(0.9, 0.999)
Batch Size	128	128	128	128	512
Warmup Steps	1000
Learning Rate Schedule	Cosine Decay
Weight Decay	1e-6
Dropout	0.1	0.1	0.1	0.1	0.0
Gradient Clipping	1.0	1.0	1.0	1.0	0.2
EMA Decay	0.9999
Training Steps	250k	250k	250k	250k	500k
Max Seq Length (Source)	n/a	n/a	256	64	128
Training Time (BART-base)	1d 11h	1d 20h	2d 22h	1d 20h	—
Training Time (FLAN-T5-base)	1d 17h	1d 21h	4d 2h	2d 7h	—
Training Time (MT5-base)	—	—	—	—	9d 16h
Table 15: Training details for our autoregressive baseline across different datasets.
ROCStories	AG News
Model	GPT-2-Medium
Trainable Params	355M
Max Seq Length	64
Optimizer	AdamW [39]
Learning Rate	8e-5
(β1, β2)	(0.9, 0.999)
Batch Size	32
Warmup Steps	500
Learning Rate Schedule	Linear Decay
Weight Decay	1e-2
Dropout	0.1
Gradient Clipping	1.0
Training Steps	100k
Table 16: Training details for our Seq2Seq baselines model.
XSum	QQP
Model	BART-base	FLAN-T5-base	BART-base	FLAN-T5-base
Trainable Params	139M	220M	139M	220M
Max Seq Length (Source)	256	64
Max Seq Length (Target)	64
Optimizer	AdamW [39]
Learning Rate	5e-5	1e-4	5e-5	5e-5
(β1, β2)	(0.9, 0.999)
Batch Size	32
Warmup Steps	500
Learning Rate Schedule	Linear Decay
Weight Decay	1e-2
Gradient Clipping	1.0
Training Steps	100k
pronounced for the AG News dataset which has a more diverse vocabulary and uses many proper nouns such as names that are out-of-vocabulary. We also present random class-conditional samples for the AG News datasets (Table 19 and Table 20) for all of the classes.
We present examples of sequence-to-sequence generations for QQP in Table 21 and XSum in Table 22. While the DiffuSeq generations are somewhat reasonable for the simpler QQP paraphrasing dataset, the model completely fails to produce coherent summaries for the challenging XSum dataset. This is the case even though DiffuSeq is trained for significantly longer than LD4LG and uses  $8 \times$  as many sampling timesteps.