SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control

Introduction

Diffusion models (Sohl-Dickstein et al., 2015), trained to iteratively refine noised inputs, have recently emerged as powerful tools for generative modeling in several continuous-valued domains such as images (Ho et al., 2020), audio (Kong et al., 2021), video (Ho et al., 2022), among others. Attempts to adapt them for discrete domains such as text data, however, have only had limited success: prior work have shown to be promising on specialized cases and small datasets (Hoogeboom et al., 2021; Austin et al., 2021; Li et al., 2022; Chen et al., 2022), but diffusion models for text still underperform (and thus are not widely adopted) compared to autoregressive language models (AR-LMs) which remain the state-of-the-art general purpose text generators (Radford et al., 2019; Brown et al., 2020).

Despite potential advantages of diffusion models for text, there are two key challenges. First, diffusion models generate text non-autoregressively, i.e., they generate (and update) the entire sequence simultaneously rather than token by token left-to-right. Although this property is useful in practice since each output token is informed by a broader bi-directional context (Lee et al., 2018; Ghazvininejad et al., 2019), it requires pre-defining an output sequence length. This limits the flexibility and applicability of trained models. On the other hand, non-autoregressive training with long sequences is expensive and difficult to optimize. In this work, we propose a semi-autoregressive solution which strikes a balance between length flexibility and the ability to alter previously generated tokens.

A major advantage of diffusion models over the current standard of autoregressive LMs is their post-hoc controllability using guidance from auxiliary models such as style classifiers (Dhariwal and Nichol, 2021). However, controllability is hard to achieve without compromises in modularity in diffusion-based LMs for text. To enable diffusion generation into discrete text rather than continuous modalities, prior approaches have employed different approximations, e.g., training with embeddings, character, or byte-level methods (Li et al., 2022; Hoogeboom et al., 2021; Austin et al., 2021; Chen et al., 2022). In contrast, existing mainstream LMs and the guidance classifiers they derive often operate at a sub-word level with sub-word representations trained jointly with the language model (Devlin et al., 2019; Liu et al., 2019; Raffel et al., 2020). Subsequently, changing the input representations to characters or embeddings requires developing guidance models from scratch, which can be expensive or infeasible in many cases. In this work, we propose a simplex-based solution which enables the diffusion over discrete texts while maintaining the advantages of diffusion models with plug-and-control guidance models.

In sum, to enable diffusion-based LMs for text we present Ssd-LM (§3), addressing the above two challenges. Ssd-LM is trained to generate text semi-autoregressively—generating blocks of tokens left-to-right with bidirectional context within the block—which offers the benefits of both AR-LMs and diffusion models. It supports training with and generating variable-length sequences. At the same time, it allows refinement within the token block, in contrast to token-level autoregressive decoding where previously generated tokens cannot be modified at all. Ssd-LM uses the same tokenization as popular AR-LMs, representing discrete text via a distribution (or simplex) defined over the vocabulary and is trained to reconstruct texts from noisy versions of the distributions. Due to its underlying representation, our method also offers an easy and modular way of guided (controlled) generation using off-the-shelf text classifiers under the minimal assumption of shared tokenizer.

Our evaluation experiments show, for the first time, that a diffusion-based LM matches or outperforms strong AR-LMs on standard text generation benchmarks (§4). We evaluate Ssd-LM on two tasks: (1) unconstrained prompt-based generation substantially outperforming existing diffusion LM approaches and performing on par with or outperforming strong autoregressive LM GPT-2 (Radford et al., 2019) on both quality and diversity (§4.2); and (2) controlled text generation with guidance from off-the-shelf classifiers (no post-hoc training/adaptation) outperforming competitive controlled text generation baselines (§4.3).

Background

Since their inception as image generators, diffusion models (and their cousins score-based models (Song and Ermon, 2019)) have been widely adopted as high-quality generative models for multiple data modalities. Here, we briefly describe a simplified view of a canonical method, denoising diffusion probabilistic models (Ho et al., 2020, DDPM) which we adapt in this work for text generation. We assume a given dataset $\mathcal{D}=\{{}^{1}\boldsymbol{x}_{0},\ldots,{}^{N}\boldsymbol{x}_{0}\}$ of continuous valued items ${}^{i}\boldsymbol{x}_{0}$ (e.g., pixel values of an image) henceforth referred to as $\boldsymbol{x}_{0}$ for simplicity.

Training a diffusion model first involves adding a series of Gaussian noise to the original data $\boldsymbol{x}_{0}$, through $T$ timesteps:

where $t\in(1,T)$ and $\boldsymbol{\epsilon}_{t}\sim\mathcal{N}(\boldsymbol{0},\mathbf{I})$. $\bar{\alpha}_{t}=\prod_{t^{\prime}=1}^{t}\alpha_{t^{\prime}}$, where $\alpha_{t^{\prime}}$ follow a predefined schedule such that $\bar{\alpha}_{t}\to 0$ as $t\to T$. This process is called forward diffusion. A diffusion model (parameterized by $\theta$) is trained to reverse this forward process by predicting the added noise $\boldsymbol{\epsilon}_{t}$ given $\boldsymbol{x}_{t}$ with the following loss:

Inference

To get an output from this model, we sample $\boldsymbol{x}_{T}\sim\mathcal{N}(\boldsymbol{0},\mathbf{I})$ and iteratively reconstruct a sample $\boldsymbol{x}_{0}$ by going back in time,

for $t=T,\ldots,1$.We omit an additional noise term $z$ here for simplicity, which is present in DDPM but not in another variant DDIM Song et al. (2021). The key obstacle in using vanilla diffusion models directly as text generators is that language consists of discrete tokens, i.e., a non-continuous $\boldsymbol{x}_{0}$ to which a continuous valued Gaussian noise cannot be added. We propose a straightforward and effective solution by treating tokens as continuous valued simplexes over the vocabulary (Hoang et al., 2017). Other existing methods addressing this problem are discussed in §5.

2 Autoregressive LM

An autoregressive LM model optimizes for the likelihood of a sequence of tokens $w^{0},\ldots,w^{L-1}$.

To decode from AR-LMs, one can provide a context $\boldsymbol{w}^{<c}$ and decode the next token $w^{c}$ iteratively by predicting $p_{\theta}(w^{c}\mid\boldsymbol{w}^{<c})$ and sampling from it to get the discrete token (Fan et al., 2018; Holtzman et al., 2020). Prior work has shown that these decoding approaches (and by extension the LMs themselves) are prone to degrade when generating long sequences and often devolve into repeating subsequences (Holtzman et al., 2020; Meister et al., 2022). In addition, such LMs do not provide a natural way to incorporate sequence-level control as tokens are generated one at a time without the ability to modify previously generated tokens (Dathathri et al., 2020; Kumar et al., 2022b). In this work, we present a method to train a semi-autoregressive LM that decodes blocks of $B$ tokens at a time, alleviating said issues with the support of diffusion models. Existing literature addressing the two issues individually are discussed in §5.

Ssd-LM

We introduce Ssd-LM—Semi-autoregressive Simplex-based Diffusion Language Model— adapting key components from both autoregressive LM and vanilla diffusion models. Conceptually, Ssd-LM uses diffusion model to decode $\boldsymbol{w}^{c:c+B}$, a block of tokens of length $B$, given a Gaussian noise and a context $\boldsymbol{w}^{<c}$ of length $c$. We show an intuitive diagram and pseudo-code for the training and decoding algorithm of Ssd-LM in Figure 1, Figure 2, and Figure 3.

Forward diffusion

Following Ho et al. (2020), we add a time-dependent Gaussian noise to the logits.

Loss function

In Equation 2, a diffusion model is trained to predict the added noise from the noisy representations. Since the forward diffusion process can be computed in a single step (Equation 1), the notion here is equivalent to predicting the original data representation (Song et al., 2021; Li et al., 2022). Our objective follows the same intuition but estimates a likelihood instead of the L2 distance while conditioning on additional context:L2 distance did not work in our pilot study potentially due to the intrinsically skewed simplex representation.

In §A, we present another interpretation of the training objective as an intuitive contrastive loss.

2 Decoding

which are then converted to a distribution via softmax. To feed this output to the next step of reverse diffusion, $t-1$, we define a $\operatorname{logits-projection}$ operation to build a predicted data representation close to the initial data representation (almost-one-hot mapping; Eq. 5). We consider three projection operations.

Greedy: creates an almost-one-hot logit centered at the highest probability token.This shares a similar intuition as a greedy clamping trick in the embedding-based diffusion in Li et al. (2022).

Sampling: creates an almost-one-hot logit centered around a token sampled from the output distribution using top-$p$ sampling (Holtzman et al., 2020). $p$ is a hyperparameter.

Multi-hot: creates an almost-one-hot logit centered around all tokens in the top-$p$ nucleus.

Decoding iteration

for $t=T,\ldots,1$ and $\boldsymbol{z}\sim\mathcal{N}(\boldsymbol{0},K^{2}\mathbf{I})$.

It is worth noting that our proposed decoding algorithm is novel and different from the DDPM decoding (Equation 3). The DDPM decoding is designed for diffusion in a continuous space and failed to generate sensible outputs in our preliminary experiments based on simplexes. In §B, we draw a theoretical connection between our decoding algorithm and DDPM decoding, and also highlight the intuitive difference between the two.

Highly-modular control

A useful property of continuous diffusion models that naturally arises from their definition is the ability to guide the generated samples to have user-defined attributes at test time. This can be done using gradients from auxiliary models such as classifiers (Dhariwal and Nichol, 2021), e.g., guiding the output of an LM to be of a positive sentiment using a sentiment classifier. There is a vibrant community of developers on platforms such as HuggingFace where many such text classifiers are publicly available. The underlying data representation of Ssd-LM is based on vocabulary simplexes. Hence, as long as a classifier shares the same tokenizer as the LM, it can be used for control in an off-the-shelf manner without modifications. This is in contrast to prior work in diffusion language models that do not support such classifiers due to differences in their input representation space (Hoogeboom et al., 2021; Austin et al., 2021; Li et al., 2022; Chen et al., 2022) and require retraining the classifiers from scratch. This ability makes Ssd-LM highly modular for controlled text generation and offers key benefits: (1) Training accurate classifiers for many tasks requires huge amounts of data where retraining them can be quite expensive, and (2) this approach allows control from classifiers that are open to use but have been trained on closed source data.

To guide Ssd-LM to generate texts with a target attribute $y$ via a standalone attribute model $f_{\phi}(\cdot)$, we update $\boldsymbol{w}_{\text{logits},t}^{c:c+B}$ (Eq. 10) at each timestep $t$ to the form below, drifting according to the gradients from the attribute classifier.

where $\lambda$ is a hyperparameter balancing the weight of control. The parameters of the standalone attribute model $\phi$ are frozen. We make a trivial modification to the embedding computation as in §3.1, to allow the classifier to take as input a simplex.

3 Additional details

We follow Nichol and Dhariwal (2021) for a cosine schedule of $\bar{\alpha}_{t}$:

where $s$ is small offset set to 1e-4 in our work and $\alpha_{t}=\frac{\bar{\alpha}_{t}}{\bar{\alpha}_{t-1}}$.

Fewer timesteps T𝑇T in decoding

Decoding from diffusion models requires a series of timesteps ($T$) which can be computationally expensive if $T$ is large. Following Li et al. (2022), we consider using a smaller value of $T$ at test time to improve decoding speed. In this work, we primarily experiment with $T_{\text{decode}}=\frac{T_{\text{train}}}{2}$ and $T_{\text{decode}}=\frac{T_{\text{train}}}{5}$.

Flexible decoding block size B𝐵B

Our Ssd-LM is trained with a fixed token block size $B_{\text{train}}$. However, the decoding algorithm has a freedom to use a different $B_{\text{decode}}$. In our experiments, we consider both scenarios of $B_{\text{train}}=B_{\text{decode}}$ and $B_{\text{train}}\neq B_{\text{decode}}$. Nevertheless, we leave for future work a more detailed analysis of the impact of the difference between $B_{\text{train}}$ and $B_{\text{decode}}$ on model performance.

Experiments

We use a bidirectional Transformer encoder RoBERTa-large (Liu et al., 2019) (0.4B, comparable size to GPT2-medium) as Ssd-LM’s underlying architecture.We initialize the model with RoBERTa’s weights as well. We observe in our initial exploration that it helps the training loss converge faster than a randomly initialized model. However, given enough computational resources, we conjecture that a randomly initialized model will offer similar performance. Note that RoBERTa uses a general BPE tokenization (Sennrich et al., 2016), same as a variety of LMs such as GPT-2 (Radford et al., 2019), GPT-3 (Brown et al., 2020), OPT (Zhang et al., 2022), etc. Any attribute classifier using the same tokenization strategy can be used to control Ssd-LM in a highly modular way.

Pretraining data, constants, and resource

We train Ssd-LM on the same data as GPT2 to make fair comparisons possible: OpenWebText (Gokaslan and Cohen, 2019) which contains 9B tokens. Following Zhang et al. (2022), we consider this data as one contiguous sequence of tokens and break it into sequences of length 200 (same as the maximum sequence length our model accepts). We randomly sample 99% of these sequences for pretraining while leaving the rest as held out for evaluation. We use the following model hyperparameters:Future work can do a search given more resources.

We use an aggregated batch size of 6,144 and a learning rate of 1e-4 with an AdamW optimizer (Loshchilov and Hutter, 2019). We trained Ssd-LM for 100K steps, which took about 6 days on 32 Nvidia V100 GPUs.

Pretraining loss

Canonical training-time perplexity of LMs is not compatible with diffusion LMs due to the difference in the inputs to the models (Equation 4 and Eq. 9). Our pretraining loss is a per-token negative log-likelihood (NLL) that depends on the specific noise schedule being used. Ssd-LM gets an average NLL of 3.87 at the end of pretraining. We show a pretraining loss curve in the appendix (§D).

2 Unconstrained text generation

First, we benchmark Ssd-LM with autoregressive LMs trained on the same data (GPT2) on text generation quality. We randomly sample 1000 sequences from the held-out OpenWebText test data, extract their prefixes as prompts (context), and generate continuations from the LMs. We consider three setups: with prompt lengths 25, 50 and 100 with respective output lengths as 25, 50 and 100 tokens. In each setup, we sample 5 continuations for each input context, thus comparing the quality of 5,000 generations from baseline GPT-2 models and our Ssd-LM.

We compare Ssd-LM with GPT2-medium, large and xl models (containing 0.4B, 0.8B and 1.6B parameters respectively) as baselines. For reference, our model size is comparable to GPT2-medium. We experiment with two popular decoding strategies for the baseline GPT-2 models with canonical parameters: nucleus sampling (Holtzman et al., 2020) with a top-$p$ of 0.9 and 0.95, and typical sampling (Meister et al., 2022) with a typical-$\tau$ of 0.2 and 0.95.

For Ssd-LM, we consider three logits projection strategies, sampling and multi-hot with $\text{top-}p\in\{0.0,0.1,0.2,0.5,0.7,0.9,0.95,0.99\}$, and greedy (which is functionally equivalent to the sampling with top-$p$=0). We use a test block size ($B_{\text{decode}}$) of 25. When generating samples of length 50 or 100, we semi-autoregressively sample in blocks of 25 and feed them as additional context to generate the next block as described in §3.2.

We evaluate the generated continuations on two axes: quality and diversity. As automatic quality metrics, we report perplexity measured by a separate, larger language model (GPT-Neo-1.3B, Black et al., 2021). Prior works, however, have shown that low perplexity of generated text is not necessarily an indication of high quality but of degenerate behavior (Nadeem et al., 2020; Zhang et al., 2021) and have proposed closeness to the perplexity of human-written text as a better evaluation. Hence, we also report the difference of log perplexity between the generated text and human-written continuations ($|\Delta_{\log\text{PPL}}|$). For diversity evaluation, we report Zipf’s coefficient (Zipf) and average distinct $n$-grams in the output samples (Li et al., 2016, Dist-$n$). In addition, we also report the repetition rate (Welleck et al., 2020; Holtzman et al., 2020, Rep), measuring the proportion of output samples that end in repeating phrases. Finally, we report MAUVE (Pillutla et al., 2021) which evaluates both quality and diversity together by approximating information divergence between generated samples and human-written continuations (from the OpenWebText held-out set).

Results

footnote 9 summarizes our main results on the 50-token prompt and output setup. We report the numbers for the best performing three settings for logits projection and decoding steps $T$ in Ssd-LM. We report the best setting for the baselines. The results for other generation lengths have a similar trend and can be found in the appendix (§D).

We find that Ssd-LM, though being smaller in size, outperforms larger GPT-2 models on the unified metric MAUVE. On diversity, Ssd-LM outperforms GPT-2 in Dist-$n$ while achieving lower repetition rates. On perplexity, the results are slightly mixed. We observe a trade-off between MAUVE and perplexity for different settings we considered, indicating that further tuning of the hyperparameters may be required. However, one of our best performing settings (sampling top-$p$=0.9, $T$=2500) still achieves the closest perplexity to the gold continuation.

In §D, we show the influence of different logits projection strategies and the associated parameters on the output text quality in Figure 4. We also show qualitative examples of the generations by Ssd-LM in Appendix D and a trajectory of intermediate states during the decoding process in Appendix D.

Comparison with Li et al. (2022)

A prior work to us, Li et al. (2022) propose Diffusion-LM, an embedding-based diffusion model trained on two small toy datasets, E2E (Novikova et al., 2017) and ROCStories (Mostafazadeh et al., 2016). In this subsection, we make a diversion to compare the embedding-based Diffusion-LM with our semi-autoregressive, simplex-based Ssd-LM. Following Li et al. (2022), we train a Diffusion-LM on ROCStories with a default embedding size of 128, 0.1B parameters under a BERT-base (Devlin et al., 2019) structure,We train two versions of Diffusion-LM, with and without BERT’s encoder weights as an initialization. The default no-initialization setup as in Li et al. (2022) works reasonably, while the other degenerates. Details can be found in §C. and a sequence length of 100. For a fair comparison, only within this subsection we train a Ssd-LM with ROCStories sequences of 100 tokens, a decoding block size of 25, and a BERT-base initialization. Further details of the setup can be found in §C.

On 2,700 held-out ROCStories sequences, we use the first 50 tokens of each sequence as a prompt and have the model generate the next 50. In footnote 10, we show the MAUVE score and perplexity of both models. We observe a substantially higher MAUVE score and lower perplexity with Ssd-LM.

3 Controlled text generation

To evaluate Ssd-LM’s ability for highly-modular control, we consider the task of sentiment controlled generation where given a prompt, the goal is to generate a continuation with a positive (or negative) polarity. We use a set of 15 short prompts as in Dathathri et al. (2020) and generate 20 samples per prompt per sentiment category, making the total number of generated samples to be 600. Following Mireshghallah et al. (2022), we generate samples with 3 different output lengths: 12, 20 and 50. For guidance, we simply import a popular sentiment classifierhttps://huggingface.co/cardiffnlp/twitter-roberta-base-sentiment from HuggingFace trained with Twitter sentiment data with over 58M training examples (Barbieri et al., 2020). This model serves as $f_{\phi}(\cdot)$ as shown in Eq. 16. In addition to quality and diversity of the generated samples, we also evaluate them on control (that is measuring if the generated output is actually positive or negative in polarity). For this, we use an external sentiment classifier trained on a different dataset. Specifically, we use a classifier trained with Yelp reviewshttps://huggingface.co/textattack/bert-base-uncased-yelp-polarity (Zhang et al., 2015; Morris et al., 2020) following the evaluation setup in the baselines we consider.

Again, we consider the sampling and multi-hot decoding strategies with $\text{top-}p\in\{0.2,0.5,0.9\}$, $T_{\text{decode}}\in\{1000,2500,5000\}$, and the multiplier for control $\lambda\in\{0,100,500,2000\}$. For the generation of 12/20/50 tokens, we use $B_{\text{decode}}$=12/20/25 and apply the decoding algorithm for $m$=1/1/2 iterations respectively.

Results

We show the quality of the controlled generations from three perspectives: target attribute via the external classifier accuracy, fluency via perplexity, and diversity via the distinctiveness measures. In footnote 14, we show the experimental results for output length 50. The results at length 12 and 20 have a similar trend and can be found in the appendix (§D).

Among the baseline methods, DAPT (Gururangan et al., 2020), GeDi (Krause et al., 2021), and DExperts (Liu et al., 2021) require training customized language models aware of the desired attributes (denoted as CM in Table 7). PPLM (Dathathri et al., 2020), FUDGE (Yang and Klein, 2021), and MuCoLa (Kumar et al., 2022b) require training a customized attribute classifier (CC). While our proposed method Ssd-LM and M&M LM (Mireshghallah et al., 2022) can directly import mainstream existing attribute classifiers from platforms like HuggingFace and are thus highly modular (HMC). We show the baseline results as reported in Mireshghallah et al. (2022) and Kumar et al. (2022b).

Ssd-LM shows strong controllability while possessing great modularity. Ssd-LM outperforms M&M LM, the other HMC method by a large margin. Even when comparing with the CC and CM methods, our method achieves a good balance in control, fluency, and diversity.

In §D, we show the impact of the control weight $\lambda$ and top-$p$ on the attribute accuracy and perplexity in Figure 5. We also show qualitative examples of the controlled generations by Ssd-LM in Appendix D.

Related work

Diffusion models have demonstrated impressive performance in popular continuous-valued domains such as images (Ho et al., 2020), audio (Kong et al., 2021), video (Ho et al., 2022) and recently also been adopted for 3D-shapes, protein structures, and more (Zhou et al., 2021; Trippe et al., 2022; Wu et al., 2022). Since they are based on adding Gaussian noise, these approaches are not straightforward to apply to discrete valued domains like text. Hoogeboom et al. (2021); Austin et al. (2021) propose diffusing in the discrete space using categorical distributions which are modified using transition matrices. However, these methods do not straightforwardly support control and yield worse results than comparable autoregressive models. Li et al. (2022) propose to represent each token as a continuous embedding and apply diffusion in the embedding space. They train the LM to generate a fixed length sequence whereas Ssd-LM allows flexibility in the generated sequence length by generating block-wise. Further, their LM is trained with specialized datasets and not evaluated against general-purpose autoregressive LMs on unconstrained text generation. Their method supports post-hoc control but requires training a customized attribute classifier,The control for diffusion models can also be classifier-free (Ho and Salimans, 2021) but requires training with the target attribute in advance, which is not a focus of this work. since the diffusion operates on a learned embedding space. Gong et al. (2022), a concurrent work to ours, extend Li et al. (2022) to a sequence-to-sequence setup with a similar underlying embedding-based method. Our work is most closely related to Chen et al. (2022) which transform discrete data into a sequence of bits and represent each bit as +1 or -1 converting it into a continuous-valued domain. For textual data, however, it can lead to extremely long sequences which are difficult to optimize. In this work, we instead maintain a subword based vocabulary but represent each token as a sequence of manually defined logits.

Language models

The majority of existing language models for text generation are trained autoregressively, i.e., they predict the next token given previously generated context. This paradigm scaled up both in terms of model size and training data size has resulted in impressive capabilities on many benchmarks (Brown et al., 2020; Chowdhery et al., 2022). However, they generate text one token at a time which does not provide flexible control over attributes of the generated text. Non-autoregressive models which generate the entire output sequence at the same time have also been explored in prior work other than diffusion models (Lee et al., 2018; Ghazvininejad et al., 2019). However, they are primarily focused on improving decoding efficiency and applied for specialized tasks like translation (Gu et al., 2018; Kaiser et al., 2018; Wang et al., 2019) and text editing (Gu et al., 2019). Many of these work have iterative processes in a discrete space, with some exploring continuous representations (Ma et al., 2019; Lee et al., 2020). To address the quality decline with the non-autoregressive methods compared to autoregressive models, prior work have also explored semi-autoregressive approaches (Wang et al., 2018; Qi et al., 2021). In the same vein, our work seeks to address the drawbacks of autoregressive language models and non-autoregressive diffusion models with a middle ground.

Controllable text generation

Early solutions for controlling attributes of generated text focused on training or finetuning AR-LMs with specific control codes (Keskar et al., 2019; Gururangan et al., 2020; Chan et al., 2021). These methods are difficult to extend to new controls as it requires retraining the models. More recent work includes decoding approaches from pretrained AR-LMs without modifying the models, through altering the output probability distribution at each step using different control objectives (Dathathri et al., 2020; Krause et al., 2021; Yang and Klein, 2021; Liu et al., 2021; Lu et al., 2021; Pascual et al., 2021). However, these methods do not allow modifying a token once it is generated and are thus suboptimal for controls at the scope of the whole sequence. Closely related to Ssd-LM are Kumar et al. (2021); Qin et al. (2022); Kumar et al. (2022b), which propose gradient-based decoding algorithms from AR-LMs. They require computing a backward pass through the LMs for each iteration, an expensive operation. In contrast, Ssd-LM with its semi-autoregressive setup allows editing past tokens via diffusion. In addition, most of these approaches require training control functions from scratch whereas our model allows using off-the-shelf classifiers. Mireshghallah et al. (2022) propose a non-autoregressive LM based on Metropolis-Hastings sampling. It also supports off-the-shelf classifiers for control, and we therefore use it as a direct baseline for Ssd-LM.

Conclusion

We present Ssd-LM, a semi-autoregressive diffusion based language model trained to denoise corrupted simplexes over the output vocabulary. Compared to prior work in text-based diffusion, Ssd-LM offers more flexibility in output length by generating blocks of text and an ability to use off-the-shelf attribute classifiers for control without additional tuning. On unconstrained text generation, Ssd-LM performs on par with or outperforms strong and larger autoregressive baselines (GPT-2) in generation quality and diversity, while vastly outperforming diffusion baselines (Diffusion-LM). On controlled text generation, Ssd-LM surpasses baselines while possessing an easy-to-use modular design. We believe that Ssd-LM opens an exciting direction for future research in flexible and modular diffusion-based language generation.

Limitations

In AR-LMs, an NLL loss is computed at training time for every token in the sequence of length $L$ (Equation 4). However, in Ssd-LM, each time a pretraining example is sampled, the loss is computed on only $B$ tokens (Eq. 9) leading to a lower sample efficiency than AR-LM. Towards improving this efficiency, future work could explore model architectures dedicated to semi-autoregressive diffusion rather than the vanilla Transformer encoder we use in this work.

Decoding speed

Since each block is generated by refining over several iterations, Ssd-LM has a considerably slower decoding speed than autoregressive models. For example, given a context of 50 tokens (single instance, unbatched), it takes Ssd-LM 25 seconds to generate the next block of 25 tokens ($T_{\text{decode}}$=1000). While our work focused on establishing the efficacy of diffusion-based LMs and modular controlled generation, future work could explore tuning $T_{\text{decode}}$ to balance model performance and decoding speed, or more efficient training and decoding algorithms extending ideas from prior work on diffusion models for continuous domains (Song et al., 2021; Nichol and Dhariwal, 2021; Rombach et al., 2022; Meng et al., 2022).

Decoding block size

In this work, although we allow setups where $B_{\text{train}}\neq B_{\text{decode}}$, the decoding block size $B_{\text{decode}}$ remains the same across $m$ decoding iterations, leaving space for a more flexible decoding schedule. Future work can also explore learning $B_{\text{decode}}$ (and $B_{\text{train}}$) rather than using constant pre-defined lengths.

Larger scale experiments with different kinds of controls and their combinations can be done, as well as more sophisticated ways to incorporate them (Kumar et al., 2021). In addition, we plan to explore alternative methods to continuously represent and add noise to discrete text (Bakosi and Ristorcelli, 2013). This work experiments with pretraining data that is primarily in English. Future work can also explore challenges and benefits of diffusion-based LMs in a multilingual setup.

Ethics statement

Language models trained on data from the web can perpetuate social biases and toxic interactions, and can be prone to generating harmful language (Gehman et al., 2020; Wallace et al., 2019, 2020; Sheng et al., 2021; Weidinger et al., 2022). Further, language generation models could memorize and amplify patterns in data without deeper language understanding or control, so they can be factually inconsistent and generate disinformation (Maynez et al., 2020; Pagnoni et al., 2021; Zellers et al., 2019), or can compromise user privacy (Carlini et al., 2021). Prior works have outlined these risks (Sheng et al., 2021; Weidinger et al., 2021), discussed their points of origin, and advocated for future research on ethical development of LMs (Bender et al., 2021; Solaiman et al., 2019).

While these studies have been conducted for autoregressive LMs, our diffusion-based LM is subject to these problems as well. However, since our method naturally incorporates controllability, future work may explore control functions that could potentially alleviate these issues (Liu et al., 2021; Kumar et al., 2022b). One risk is that controllability can also be misused maliciously, with models being intentionally exploited to generate biased, toxic, or non-factual content (Bagdasaryan and Shmatikov, 2022; Pagnoni et al., 2022). Therefore, apart from controlled generation, future work should aim to detect the generations under control as well to defend against the malicious use (Kumar et al., 2022a).

Acknowledgements

The authors would like to thank Tianxiao Shen, Tianxing He, Jiacheng Liu, Ruiqi Zhong, Sidney Lisanza, Jacob Gershon, members of TsvetShop, and the anonymous ACL reviewers for their helpful discussions and feedback. X.H. gratefully acknowledges funding from the UW-Meta AI Mentorship program. S.K. gratefully acknowledges a Google Ph.D. Fellowship. Y.T. gratefully acknowledges an Alfred P. Sloan Foundation Fellowship. This research is supported in part by by the National Science Foundation (NSF) under Grants No. IIS2203097, IIS2125201, and NSF CAREER Grant No. IIS2142739. This research is supported in part by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (IARPA), via the HIATUS Program contract #2022-22072200004. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies, either expressed or implied, of ODNI, IARPA, or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for governmental purposes notwithstanding any copyright annotation therein.

References

Appendix A A contrastive interpretation of the training loss

Optimizing $\theta$ is a contrastive objective: maximizing the estimated likelihood of true data, while penalizing the estimated likelihood of noisy data under a broad range of different noise scales.

Appendix B Connection between our decoding algorithm and the DDPM decoding

We revisit the decoding step in DDPM introduced in Equation 3. Since we know that during the training phase $\boldsymbol{x}_{t}$ is generated through a one-step forward diffusion process (Equation 1), a model $\theta$ predicting the added noise $\epsilon_{\theta}(\boldsymbol{x}_{t},t)$ can therefore be considered as predicting an imaginary $\boldsymbol{x}_{0}$ in one-step:

Below we write $\hat{\boldsymbol{x}}_{0}(\boldsymbol{x}_{t},t,\theta)$ as $\hat{\boldsymbol{x}}_{0}$ and $\epsilon_{\theta}(\boldsymbol{x}_{t},t)$ as $\epsilon_{\theta}$ for simplicity.

Rearranging the DDPM decoding transition (Equation 3), we have:

with $\sqrt{\frac{\alpha_{t}-\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}}}\approx 1$ for most $t\in(1,T)$.Specifically, we adopt a cosine schedule for $\bar{\alpha}_{t}$ (Nichol and Dhariwal, 2021), and $\sqrt{\frac{\alpha_{t}-\bar{\alpha}_{t}}{1-\bar{\alpha}_{t}}}>0.98$ for 98% of all $t$, with some outliers as $t\to 0$ and $t\to T$.

Noting the format simlarity between Equation 1 and Eq. 23, we therefore interpret the DDPM decoding transition from $\boldsymbol{x}_{t}$ to $\boldsymbol{x}_{t-1}$ as (1) predicting an imaginary $\hat{\boldsymbol{x}}_{0}$, and (2) applying a compensating forward diffusion step with a deterministic noise $\epsilon_{\theta}$.

Our decoding strategy in Eq. 15 is in a very similar form as Eq. 23. We also predict the initial data representation with $\theta$ and apply a forward diffusion step. The difference is that we sample a noise $\boldsymbol{z}$ instead of using the deterministic $\epsilon_{\theta}$, to encourage exploration.

Appendix C Detailed setup of the comparison with Diffusion-LM (Li et al., 2022)

We apply block concatenation on ROCStories similarly as OpenWebText, resulting in 50K training sequences of 100 tokens. We train Diffusion-LM with a default batch size of 64, learning rate of 1e-4, and 400K steps. We train Ssd-LM with a batch size of 512, learning rate of 1e-4, and 20K steps. Both models use a tokenizer of BERT-base-uncased. For Ssd-LM, additional hyperparameters like decoding block size and one-hot constant remain the same as the main Ssd-LM benchmarked with GPT-2. For Diffusion-LM, the evaluation in the main paper is an infilling task. We use same decoding hyperparameters as Li et al. (2022). For Ssd-LM, the evaluation is a block-wise generation problem with $m$=2 iterations. The result of Ssd-LM in footnote 10 is obtained with a decoding configuration of $T_{\text{decode}}$=2500 and top-$p$=0.5.

Our Ssd-LM in this subsection is initialized with BERT. For a fair comparison, apart from the default Diffusion-LM reported in footnote 10, we train another Diffusion-LM initialized with the encoder weights of BERT. However, this leads to degenerated results that are much worse than the default Diffusion-LM and our Ssd-LM: a MAUVE score of 0.4 out of 100 and a PPL of 73157. This problem is not due to overfitting, as all checkpoints of the model show the same degenerated result. Since Li et al. (2022) did not explore this setup in their original work as well, we conjecture that Diffusion-LM may be incompatible with pretrained weights from existing non-diffusion models by nature, a disadvantage to our Ssd-LM.

Appendix D Additional results

Figure 4 shows the influence of different logits projection strategies and the associated parameters on the unconstrained generations’ output text quality. We observe that reducing top-$p$ $\to$ 0 (greedy projection) can lead to a low perplexity but it is undesirable due to a high repetition rate. We also find the multi-hot projection strategy is overall worse performing than the sampling projection strategy in our setup, indicating it is better to commit the intermediate states to single rather than multiple tokens. This can be because our logits mapping involves putting probability mass on singular tokens. The multi-hot projection may still be a viable strategy if future work uses multi-hot logits mapping for the input tokens.

Figure 5 shows the impact of the control weight $\lambda$ and top-$p$ on the attribute accuracy and perplexity in controlled text generation. As expected, a larger control weight leads to a better external classifier accuracy. The perplexity at the same time increases with a larger $\lambda$, but under a reasonable range for a top-$p$ of 0.2 and 0.5.

Figure 6 shows the pretraining loss trajectory. Table 4, Table 5, Table 6, and Table 7 show additional evaluation results of Ssd-LM generations. Appendix D and Appendix D show qualitative examples of Ssd-LM generations.