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

Abstract: Despite the growing success of diffusion models in continuous-valued domains (e.g., images), similar efforts for discrete domains such as text have yet to match the performance of autoregressive language models. In this work, we present SSD-LM—a diffusion-based language model with two key design choices. First, SSD-LM is semi-autoregressive, iteratively generating blocks of text, allowing for flexible output length at decoding time while enabling local bidirectional context updates. Second, it is simplex-based, performing diffusion on the natural vocabulary space rather than a learned latent space, allowing us to incorporate classifier guidance and modular control using off-the-shelf classifiers without any adaptation. We evaluate SSD-LM on unconstrained text generation benchmarks, and show that it matches or outperforms strong autoregressive GPT-2 models across standard quality and diversity metrics, while vastly outperforming diffusion-based baselines. On controlled text generation, SSD-LM also outperforms competitive baselines, with an extra advantage in modularity.

## introduction diffusion models @xcite , trained to iteratively refine noised inputs, have recently emerged as powerful tools for generative modeling in several continuous-valued domains such as images @xcite , audio @xcite , video @xcite , 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 @xcite @xcite , but diffusion models for text still underperform (and thus are not widely adopted) compared to autoregressive language models (ar-lms) which 1 our code and models can be found at https:// github.com/xhan77/ssd-lm . 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-toright. although this property is useful in practice since each output token is informed by a broader bi-directional context @xcite , 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 @xcite . 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 @xcite @xcite . 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 @xcite @xcite . 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-andcontrol 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 @xcite on both quality and diversity ( §4.2); and (2) controlled text generation with guidance from off-the-shelf classifiers (no posthoc training/adaptation) outperforming competitive controlled text generation baselines ( §4.3). ## 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 w c:c+b , a block of tokens of length b, given a gaussian noise and a context @xmath0. 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 . ## related work diffusion models diffusion models have demonstrated impressive performance in popular continuous-valued domains such as images @xcite , audio @xcite , video @xcite and recently also been adopted for 3d-shapes, protein structures, and more @xcite @xcite . since they are based on adding gaussian noise, these approaches are not straightforward to apply to discrete valued domains like text. @xcite ; @xcite 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. @xcite 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, 15 since the diffusion operates on a learned embedding space. @xcite , a concurrent work to ours, extend @xcite to a sequence-to-sequence setup with a similar underlying embedding-based method. our work is most closely related to @xcite 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 @xcite . however, they generate text one token at a time which does not provide flexible control over attributes of the generated text. nonautoregressive models which generate the entire output sequence at the same time have also been explored in prior work other than diffusion models @xcite . however, they are primarily focused on improving decoding efficiency and applied for specialized tasks like translation @xcite @xcite ) and text editing @xcite . many of these work have iterative processes in a discrete space, with some exploring continuous representations @xcite . to address the quality decline with the non-autoregressive methods compared to autoregressive models, prior work have also explored semi-autoregressive approaches @xcite . in the same vein, our work seeks to address the drawbacks of autoregressive language models and non-autoregressive diffusion models 15 the control for diffusion models can also be classifierfree @xcite but requires training with the target attribute in advance, which is not a focus of this work. controllable text generation early solutions for controlling attributes of generated text focused on training or finetuning ar-lms with specific control codes @xcite @xcite . 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 @xcite @xcite @xcite . 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 @xcite ; @xcite ; @xcite , 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. @xcite 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.