Diffusion Guided Language Modeling (2024)

Abstract: Current language models demonstrate remarkable proficiency in text generation. However, for many applications it is desirable to control attributes, such as sentiment, or toxicity, of the generated language—ideally tailored towards each specific use case and target audience. For auto-regressive language models, existing guidance methods are prone to decoding errors that cascade during generation and degrade performance. In contrast, text diffusion models can easily be guided with, for example, a simple linear sentiment classifier—however they do suffer from significantly higher perplexity than auto-regressive alternatives. In this paper we use a guided diffusion model to produce a latent proposal that steers an auto-regressive language model to generate text with desired properties. Our model inherits the unmatched fluency of the auto-regressive approach and the plug-and-play flexibility of diffusion. We show that it outperforms previous plug-and-play guidance methods across a wide range of benchmark data sets. Further, controlling a new attribute in our framework is reduced to training a single logistic regression classifier.

## introduction the rapid and ubiquitous adoption of (large) language models (lms) raises a critical parallel challenge: how do we effectively guide their generation to be safe and fitting for each application and target audience? for example, one might want an lm to use different language if it interacts with kindergarteners, writes a comedy sketch, provides legal support, or summarizes news articles. currently, the most successful llm paradigm is to train a single large auto-regressive model that can be used for many tasks @xcite . different approaches to guide * correspondence @xmath0.edu>. generation of such llms exist, each with their own strengths and weaknesses. a popular way to control the generation is to align the lm through fine-tuning. these approaches are very effective, but as they change the actual model weights, they can deteriorate the lm's performance @xcite @xcite . further, if new applications require a unique combination of attribute preferences (e.g. humorous but not toxic), new dedicated models must be fine-tuned and hosted. in contrast, plug and play approaches do not change the model weights and instead utilize additional light-weight classifiers or heuristics to influence the generation process @xcite @xcite @xcite . such approaches are highly flexible and do not require fine-tuning or the hosting of dedicated models. however, as they typically alter the logits in the final layer, they are prone to creating decoding errors that cascade through the auto-regressive generation process and deteriorate the output quality. one alternative to auto-regressive generation is provided by diffusion models @xcite @xcite . originally gaining prominence in image generation, diffusion models learn to iteratively "denoise" samples of gaussian noise into samples from a target data distribution (e.g. natural images, or text completions). crucially, this iterative generation naturally allows for plug-and-play control through a simple likelihood function @xcite . minor errors introduced by the guidance mechanism can be corrected by the diffusion model later in the generative process. pre-trained image diffusion models, for instance, can incorporate plug-andplay guidance at inference-time to perform tasks such as super-resolution and in-painting, without any task-specific training. recent work has begun to explore the application of diffusion to the discrete problem of language generation @xcite @xcite @xcite . diffusion language models have demonstrated positive results in controllable generation, but still exhibit poor perplexity and generation quality compared to auto-regressive models. in this paper we propose a novel framework, diffusion guided language modeling (dglm), (see figure 1 ) that integrates the fluency of autoregressive generation with the flexibility of continuous diffusion. we develop a diffusion network that, given some text prefix, generates continuous semantic proposals of language continuations. these semantic proposals act as soft prompts and guide a fluent auto-regressive model to generate language aligned with the proposal. during pre-training, we condition the language decoder on embedded representations of the ground truth continuation, teaching the decoder that the semantic proposals contain valuable information. during inference time, we let the diffusion model generate its own proposal continuation from the prefix, guided by a simple linear classifier to ensure the desired attributes. the proposal vectors function as additional prompts for the decoder and steer it towards a fluent continuation that inherits the attributes of the proposal. dglm has several compelling properties: 1. it ## background: diffusion models we introduce diffusion models @xcite @xcite , following the presentation of @xcite most closely. given some dataset drawn from an unknown distribution q(x), our goal is to learn a generative model p θ (x), shorthanded as p(x), that approximates the unknown data distribution q(x). the observed data x could be an image, text, or some latent feature vector @xcite . forward process. diffusion models consist of a forward process and a generative process. the forward process defines a gradual transition from the data distribution to a guassian distribution. this introduces a series of increasingly noisy latent variables z t for timesteps @xmath1, 1] @xcite . this gaussian diffusion process defines the conditional distribution q(z 0,...,1 |x). for every @xmath2, 1], the marginal q(z t |x) is given by: z @xmath3, where ϵ ∼ n (0, i) the noise level is also commonly written in terms of the log signal-to-noise ratio (snr), λ @xmath4. the noise schedule, specified by α @xmath5, 1], is a strictly monotonically decreasing function defined so that the process starts with the original input, z 0 ≈ x, and the final latent becomes approximately gaussian, q(z 1 ) ≈ n (z 1 ; 0, i). generative model. the generative process reverses the forward process, defining a gradual transition from gaussian noise to the data distribution. the generative model defines a probability distribution over the latent variables, p(z 0 , . . . , z 1 ). given access to the score function ∇ z log q t (z), the gradient of the log probability density function, the forward process can be reversed exactly. diffusion models learn to approximate the score function with a neural network, s θ (z; λ)) ≈ ∇ z log q t (z), and use the estimated score function to approximately reverse the forward process. if s θ (z; λ) ≈ ∇ z log q t (z), then our generative distribution is close to the true distribution. this enables us to draw samples from a guassian distribution z 1 ∼ p(z 1 ), and approximately solve the reverse diffusion process using the estimated score s θ (z; λ). in this work, we use the standard ddpm sampler from @xcite . training objective. @xcite showed that score networks, s θ (z; λ), can be learned with a denoising score matching (dsm) loss over all data points x ∼ d and noise levels: the neural network can be parameterized in terms of the noise (ϵ), the data (x), or the velocity (v := α t x + σ t ϵ) @xcite because of the following relationship: in practice, people have found that parameterizing the neural network as an ϵ-prediction or a vprediction model improves training stability and downstream performance @xcite . we follow the best practices established in recent work @xcite and adopt the v-parameterization: the above relationships also mean that at every timestep, t, the diffusion network provides us with the minimum mean-squared error (mmse) estimate of the clean data: plug-and-play control. when drawing samples x, we want them to meet certain conditional criteria y such as a class label. one can learn the conditional score function ∇ zt log p t (z t |y) directly with a conditional diffusion model. however, like learning a conditional auto-regressive model, this would require a large corpus of annotated data and the given some prefix, we first generate an embedded representation of the language continuation with a diffusion model. during this stage, we can optionally intervene with a lightweight classifier for plug-and-play guidance. we map the continuation embedding to a soft prompt to guide an auto-regressive decoder to generate language aligned with the semantics of the generated embedding. this decomposition shows that we can perform conditional generation with an unconditional diffusion model if we can estimate ∇ zt log p t (y|z t ), the gradient of the log-likelihood of the condition given the latent @xcite . diffusion posterior sampling (dps) @xcite utilizes a conditional distribution over noiseless data p(y|x) and the mmse estimator xθ (z t , λ t ) to approximate the conditional: if the distribution of noiseless data p(y|x) is differentiable with respect to x, the dps approximation is differentiable with respect to z t . we can therefore utilize a lightweight classifier over clean data to guide an unconditional diffusion model to sample data x consistent with some criteria y in a plug-and-play manner. ## diffusion guided language modeling we present an overview of our framework in figure 2. our method has three main components-a diffusion network, a lightweight prompt generator, and a pre-trained auto-regressive decoder. given some textual prefix, we use the diffusion model to sample an embedded, semantic proposal of a possible continuation. during sampling, we can optionally perform plug-and-play control to enforce some condition (e.g. low toxicity). after sampling the semantic embedding, the prompt generator is used to process the embedding into a soft prompt, which then guides the auto-regressive decoder to generate text aligned with the proposal. ## datasets and metrics datasets. we extract a subset of 10 million instances from c4 @xcite to pre-train dglm. this represents only 2.5% of c4 and scaling the pre-training corpus would likely be fruitful. we follow @xcite and filter out uncompressible text to improve quality. if the number of gpt-2 tokens is more than @xmath6.3 times the raw number of characters, we drop it from the dataset. this removes instances consisting of, for instance, long html strings or markdown code. to evaluate the language generation capabilities of our dglm, we extract 5000 random validation instances from c4 @xcite and open-webtext @xcite . we condition the network on the first 32 tokens and generate a 32 token continuation. for our toxicity mitigation experiments, we train our logistic regression model on the jigsaw unintended bias dataset (cjadams, 2019) and evaluate the effectiveness of toxicity mitigation experiments using 5,000 neutral promps from realtoxicityprompts @xcite . for our sentiment control experiments, we utilize amazon polarity foot_0 and sst-2 @xcite to train a sentiment classifier, and perform sentiment control using 5,000 neutral prompts from openwebtext. metrics. we evaluate the fluency of text by measuring its perplexity with the open-source olmo-1b foot_1 language model. we also report mauve scores @xcite , a text generation metric that measures the similarity of generated text with that of reference text using divergence frontiers. to get embeddings for mauve, we follow the advice of @xcite and utilize electra-large @xcite . to evaluate generation diversity, we use the metric introduced by @xcite : for the guidance tasks, we generate 25 samples per prompt. we report the olmo-1b perplexity of the continuations to evaluate the fluency of the generations. we follow prior work and measure the average number of unique 3-grams, denoted dist-3, in each set of continuations to quantify generation diversity. along with ensuring that guidance does not degrade language quality or sacrifice diversity, we measure the adherence to the guidance conditions. following prior work @xcite , we use the perspective api to measure the toxicity of generated text. because @xcite found that the perspective api changes significantly over time, we re-score the released generations for all of the baselines with the current version of the api. we measure the average max toxicity across 25 generations and the toxicity rate, defined as the empirical odds of at least 1 of 25 continuations being classified as toxic. to evaluate sentiment, we utilize robertalarge 4 (liu et al., 2020 ## experimental results language generation. we validate the effectiveness of our framework on open-ended language generation in plug-and-play control. we utilize dglm to avoid generating toxic language. we show quantitative results in figure 5 and figure 7 . qualitative examples are presented in we also employ dglm to steer the sentiment of generated text. we present results for negative sentiment in figure 6 and positive sentiment in figure 8 . we observe that our method is similarly effective in this setting, decreasing (or increasing) sentiment with no loss of fluency and minimal loss of diversity for modest guidance values. compositional control. we present qualitative results from composing multiple attribute classifiers with dglm. we fine-tune an additional logistic regression model on the ag news topic classification dataset. we then sum the losses for the sentiment and topic classification classifier to guide generation. we find that dglm successfully enables compositional control and present qualitative examples in ## conclusion we present diffusion guided language modeling (dglm), a powerful integration of auto-regression and diffusion that enables versatile attribute-guided text generation with lightweight classifiers. the diffusion model generates controllable semantic proposals that guide the language decoder. extending dglm to control an unseen attribute only requires learning a single logistic regression model. experimental results show that dglm significantly outperforms prior plug-and-play methods, opening avenues for building highly adaptable lms with user-controllable behavior.