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-andplay 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. Our code is available at https://github.com/ justinlovelace/Diffusion-Guided-LM.

heading: 1 Introduction
text: 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 (Raffel et al., 2020; Brown et al., 2020). Different approaches to guide
*Correspondence to <jl3353@cornell.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 (Lazaridou et al., 2020; Ouyang et al., 2022; Bubeck et al., 2023; Noukhovitch et al., 2023). 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 (Dathathri et al., 2019; Yang and Klein, 2021; Krause et al., 2021; Liu et al., 2021; Deng and Raffel, 2023). Such approaches are highly flexible and do not require fine-tuning
14936
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 (Sohl-Dickstein et al., 2015; Song et al., 2020; Ho et al., 2020). 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 (Dhariwal and Nichol, 2021). 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 (Li et al., 2022; Gong et al., 2022; Lovelace et al., 2023; Gulrajani and Hashimoto, 2023; Zhang et al., 2023). 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
decouples model training from attribute control. 2. Controlling a new attribute only requires the training of a simple logistic regression classifier. 3. Empirically, DGLM is extremely effective and outperforms the current state-of-the-art in plug-andplay control across diverse benchmark data sets.

heading: 2 Background: Diffusion Models
text: We introduce diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020), following the presentation of Kingma and Gao (2023) 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 (Rombach et al., 2021).
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 zt for timesteps t ∈ [0, 1] (Kingma et al., 2021). This Gaussian diffusion process defines the conditional distribution q(z0,...,1|x). For every t ∈ [0, 1], the marginal q(zt|x) is given by:
zt = αtx+ σtϵ, where ϵ ∼ N (0, I)
We utilize the common variance-preserving formulation, where σ2t = 1− α2t . The noise level is also commonly written in terms of the log Signal-toNoise Ratio (SNR), λt = logα2t /σ 2 t . The noise schedule, specified by αt ∈ [0, 1], is a strictly monotonically decreasing function defined so that the process starts with the original input, z0 ≈ x, and the final latent becomes approximately Gaussian, q(z1) ≈ N (z1;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(z0, . . . , z1).
Given access to the score function ∇z log qt(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 qt(z), and use the estimated score function to approximately reverse the forward process. If sθ(z;λ) ≈
∇z log qt(z), then our generative distribution is close to the true distribution.
This enables us to draw samples from a Guassian distribution z1 ∼ p(z1), and approximately solve the reverse diffusion process using the estimated score sθ(z;λ). In this work, we use the standard DDPM sampler from Ho et al. (2020).
Training objective. Song and Ermon (2019) 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:
LDSM(x) = Et,x,ϵ[w(λt) · ∥sθ(zt;λ)−∇zt log q(zt|x)∥22],
where w(λt) is a SNR-dependent weighting term that is tuned to emphasize noise levels important for downstream sample quality.
The neural network can be parameterized in terms of the noise (ϵ), the data (x), or the velocity (v := αtx+σtϵ) (Salimans and Ho, 2022) because of the following relationship:
∇zt log q(zt|x) = −ϵ/σt = −σ−2t (zt − αtx) = −zt − (αt/σt)v.
In practice, people have found that parameterizing the neural network as an ϵ-prediction or a vprediction model improves training stability and downstream performance (Ho et al., 2020; Salimans and Ho, 2022). We follow the best practices established in recent work (Kingma and Gao, 2023) and adopt the v-parameterization:
Lv(x) = Et,x,ϵ[w(λt) · ∥v̂θ(zt;λ)− vt∥22].
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:
x̂θ(zt, λt) = αtzt − σtv̂θ(zt;λ).
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 pt(zt|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
conditional model could not be easily adapted to other conditions. We can instead use Bayes’ rule to decompose the conditional score at time t into the unconditional score and a likelihood term:
∇zt log pt(zt|y) = ∇zt log pt(zt) +∇zt log pt(y|zt).
This decomposition shows that we can perform conditional generation with an unconditional diffusion model if we can estimate ∇zt log pt(y|zt), the gradient of the log-likelihood of the condition given the latent (Dhariwal and Nichol, 2021).
Diffusion Posterior Sampling (DPS) (Chung et al., 2023) utilizes a conditional distribution over noiseless data p(y|x) and the MMSE estimator x̂θ(zt, λt) to approximate the conditional:
∇zt log pt(y|zt) ≈ ∇zt log p(y|x̂θ(zt, λt)). If the distribution of noiseless data p(y|x) is differentiable with respect to x, the DPS approximation is differentiable with respect to zt. 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.
In practice, people often introduce some guidance weight term s as a hyperparameter
∇zt log pt(zt|y) = ∇zt log pt(zt) + s · ∇zt log pt(y|zt),
where setting s > 1.0 increases the influence of the conditional information. This can be viewed as sampling from a modified distribution p̃t(zt|y) ∝ pt(zt)pt(y|zt)s.

heading: 3 Diffusion Guided Language Modeling
text: 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.

heading: 3.1 Semantic Proposal Conditioning
text: Sentence-T5 (Ni et al., 2022) is a sentence encoder that is trained contrastively, producing embeddings that capture high-level semantics while being robust to shallow surface-form variations. Because of these properties, we learn our diffusion model in its latent space to generate semantic proposals1.
In order to condition the auto-regressive decoder on embeddings from Sentence-T5, we introduce a lightweight prompt generator that maps the embedding to a soft prompt for the decoder (see Figure 3). We fine-tune the prompt generator and decoder to generate continuations that correspond to the embeddings from the frozen Sentence-T5 encoder.
Given some text sequence, we split it to obtain a prefix and continuation. We use Sentence-T5 to embed the continuation into a 768-dimensional vector, denoted xcont. The prompt generator linearly projects the embedding to dimension 4d, splits it into k = 8 feature vectors, and then further refines them using a small transformer (Morris et al.,
1We use Sentence-T5-XL in this work.
2023). This yields a sequence of k soft tokens that guide the auto-regressive model to reconstruct the continuation. The input training sequence therefore consists of the prefix text and the soft prompt, which are used to predict the text continuation with teacher forcing.
The auto-regressive model is trained with the standard language modeling loss. We mask out the predictions corresponding to the soft tokens from the loss function. Because the sentence embedding corresponds to the ground-truth continuation, the auto-regressive network will learn to generate text aligned with the Sentence-T5-XL embedding.
Gaussian noise conditioning. During generation, we will be utilizing latent proposals from our diffusion network. While an effective diffusion model produces high-quality proposals, it is difficult to match the quality of the ground-truth embeddings used during pre-training. To improve the robustness of the auto-regressive decoder to minor errors introduced by the diffusion network, we incorporate Gaussian noise augmentation, a technique introduced for cascaded image diffusion models (Ho et al., 2022; Saharia et al., 2022).
The prompt generator receives a latent sampled from the forward diffusion process zt = αtxcont + σtϵ, where the noise level is sampled according to some schedule αt ∈ [0, 1]. We also condition the prompt generator on the level of noise. The noise level dynamically adjusts the influence of the proposal embedding on the auto-regressive decoder’s output. At low noise levels the decoder relies heavily on the proposal embedding, while at high noise levels, the decoder falls back to standard auto-regressive generation.
During generation, we pass a proposal embedding with some low, but non-negligible, level of noise (we set σ2t = 0.05 by default) and the autoregressive decoder will generate text aligned with the proposal while correcting for minor errors introduced by the diffusion network. This also provides us with a knob to tailor the influence of the diffusion network to the application. We report full implementation details in Table 7.

heading: 3.2 Semantic Diffusion
text: Our semantic diffusion model operates in the latent space of Sentence-T5, iteratively generating potential text continuations guided by a text prefix. Given a text sequence, we split it into a prefix and a continuation and embed both using Sentence-T5,
denoted as xpref and xcont respectively. We train the score network to recover the noisy continuation embedding given the prefix embedding. More formally, the noisy latent is given as zt = αtxcont + σtϵ and we parameterize our score network as sθ(zt;λ;xpref). We therefore learn to sample from the distribution of possible continuation embeddings for the text prefix.
For the diffusion network, we employ a transformer model (see Figure 4). To prepare the input, we first independently project the noisy latent and prefix embeddings, then split each into 64 feature vectors. We concatenate these element-wise along the feature dimension, giving us 64 representations that we then process with the transformer.
We convert the transformer’s output to a single feature vector by inverting the initial projection operation. We down-project and concatenate the 64 feature vectors to create the final vector used for score regression. During training, we mask the prompt embedding (p = 0.1), by replacing it with a learnable null embedding, for classifier-free guidance (Ho and Salimans, 2022). This jointly trains unconditional (sθ(zt;λt)) and conditional (sθ(zt;λt;xpref)) diffusion networks. During sampling, we can use guidance weight w to blend predictions as
s̃t = wŝθ(zt;λt;xpref) + (1− w)ŝθ(zt;λt).
Setting w = 1.0 yields the conditional model while setting w > 1.0 strengthens the influence of the conditioning information and emphasizes promptadherent continuations. We report full implementation details of our diffusion model in Table 8.

heading: 3.3 Plug and Play Control
text: To effectively control text generation with desired conditions (denoted as y), we develop a plug-andplay approach leveraging the semantic structure of Sentence-T5’s embeddings. We now present the mathematical formulation of our approach.
Our semantic diffusion model estimates the score of possible text continuations within the Sentence-T5 latent space given some prefix: ∇zt log pt(zt|xpref). Given some condition y that we wish to enforce for our sample xcont at inference time, we decompose the conditional score using Bayes’ rule and the DPS approximation, x̂θ(zt, λt,xpref), as
∇zt log pt(zt|xpref,y) = ∇zt log pt(zt|xpref) +∇zt log pt(y|zt,xpref) ≈ ∇zt log pt(zt|xpref) +∇zt log p(y|x̂θ(zt, λt,xpref),xpref).
Since y depends solely on the continuation and the DPS estimate already incorporates information from the prefix, we assume conditional independence between y and xpref given x̂θ(zt, λt,xpref). Mathematically, this is expressed as:
∇zt log p(y|x̂θ(zt, λt,xpref),xpref) ≈ ∇zt log p(y|x̂θ(zt, λt,xpref)).
This simplification allows us to express the conditional score function:
∇zt log pt(y|xpref, zt) ≈ ∇zt log p(y|x̂θ(zt, λt,xpref)) = −∇ztℓy(x̂θ(zt, λt,xpref)
where ℓy(x̂θ(zt, λt,xpref) is the cross-entropy loss. With this, plug-and-play guidance simply requires a classifier within the sentence-T5 latent space. We employ a linear probe (i.e. logistic regression) in our experiments (see Appendix C for additional details). We find that semantic diffusion enables effective control with surprisingly simple classifiers.
Song et al. (2023) observed that the MMSE estimate, x̂θ(zt, λt,xpref), introduced approximation errors in the conditional score estimate. They propose sampling around the MMSE estimate
x̂(i) ∼ N (x̂θ(zt, λt,xpref), σ2t /α2t I).
The sampling distribution has large variance early in the sampling process when the DPS estimate is
uncertain and converges to the DPS point estimate at the end of the sampling process. They use a Monte-Carlo approach to approximate the guidance with the logmeanexp operation. Adapting this, we compute the guidance term as:
−∇zt log( 1
n
n∑
i
exp(ℓy(x̂ (i)))).
Early in the sampling process, this steers generation towards a region of low loss within the latent space. With n = 32, using the Monte-Carlo estimate incurs negligible overhead, requiring only 32 logistic regression evaluations.

heading: 4 Datasets and Metrics
text: Datasets. We extract a subset of 10 million instances from C4 (Raffel et al., 2019) to pre-train DGLM. This represents only 2.5% of C4 and scaling the pre-training corpus would likely be fruitful. We follow Geiping and Goldstein (2023) and filter out uncompressible text to improve quality. If the number of GPT-2 tokens is more than t = 0.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 (Raffel et al., 2019) and OpenWebText (Gokaslan and Cohen, 2019). 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 (Gehman et al.). For our sentiment control experiments, we utilize Amazon
Polarity 2 and SST-2 (Socher et al., 2013) 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-1B3 language model. We also report MAUVE scores (Pillutla et al., 2021), 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 He et al. (2022) and utilize ELECTRA-large (Clark et al., 2020). To evaluate generation diversity, we use the metric introduced by Su et al. (2022): Div = ∏4 n=2 |unique n-grams({wi})| |total n-grams({wi})| where {wi} is a set of generated samples. 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 (Deng and Raffel, 2023; Liu et al., 2021), we use the Perspective API to measure the toxicity of generated text. Because Pozzobon et al. (2023) 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 RoBERTa2https://huggingface.co/datasets/amazon_
polarity 3https://huggingface.co/allenai/OLMo-1B
Large4 (Liu et al., 2020) fine-tuned on sentiment classification across diverse domains as well as the fine-tuned DistilBERT model (Sanh et al., 2019) used by prior work.

heading: 5 Experimental Results
text: Language Generation. We validate the effectiveness of our framework on open-ended language generation in Table 1 without any plug and play control. We observe that our method achieves strong language generation results, matching or surpassing the reference perplexity with sufficient classifier-free guidance strength. We observe that DGLM leads to consistently more diverse generations than the auto-regressive baseline across both datasets. We observed that a handful of very high perplexity samples skews the GPT-2 baseline’s perplexity on C4. However, DGLM also achieves stronger MAUVE scores on that dataset.
We examine the impact of Gaussian noise augmentation in Table 2. As an additional metric, we re-embed the generated text with Sentence-T5 and compute the cosine similarity with the proposal embedding5. We observe that the Gaussian noise augmentation enables the network to smoothly interpolate between auto-regressive generation (low perplexity but poor diversity) and diffusion-guided generation (higher perplexity and diversity). We observe that lower levels of noise montonically improve the decoders adherence to the proposal.
Plug-and-Play Control. We utilize DGLM to avoid generating toxic language. We show quantitative results in Figure 5 and Figure 7. Qualitative
4https://huggingface.co/siebert/ sentiment-roberta-large-english
5We follow Zhang* et al. (2020) and rescale the cosine similarity with a baseline computed between random dataset samples.
examples are presented in Table 9. Plug-and-play guidance with a linear probe effectively mitigates toxicity with negligible trade-offs in fluency. We simultaneously achieve lower perplexity, lower toxicity, and higher diversity than all baselines.
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 Table 3 (additional examples are in Table 10). We leave rigorous evaluations of the compositionality of DGLM to future work.
Decoding Overhead. Plug-and-play methods for auto-regressive generation often introduce overhead at each decoding step. For example, DExperts employs auxiliary language models that work alongside the primary model. In contrast, DGLM incurs a one-time cost for generating the semantic proposal, which is then amortized across subsequent decoding steps. We therefore compute runtimes across a range of generation lengths. We report the relative increase in runtime compared to the original GPT-2 model for each method (baseline data from Liu et al. (2021)) in Table 4. As seen in the table, DGLM incurs a large cost for short sequences but has reduced overhead compared to prior methods with modest generation lengths.

heading: 6 Related Work
text: Fine-tuning. Continual pre-training on text from some target domain (domain-adaptive pretraining or DAPT) is an effective technique for controlling attributes in generated text (Gururangan et al., 2020). Lu et al. (2022) optimize a reward function by fine-tuning an LM with control tokens for different reward quantiles. Reinforcement Learning with Human Feedback (RLHF) involves training a reward model on human preference data that is then used to fine-tune the LM (Wu et al., 2023; Ouyang et al., 2022). Jang et al. (2023) train multiple personalized RLFH models and show that these personalized models can be used alone or in conjunction with one another to produce text with desired attributes.
Guided Generation. Finetuning LMs is expensive and therefore to reduce cost, Dathathri et al. (2019) proposed Plug and Play Language Model (PPLM), a method that used light-weight classifiers to guide frozen language models during text generation. Similarly, FUDGE (Yang and Klein, 2021) trains classifiers on partial sequences to predict whether a particular attribute is satisfied and updates the output probability distribution accordingly. Instead of using a classifier, GeDi (Krause et al., 2021) trains a small class-conditional language model to act as a discriminator and guide
the language generation. Similarly, DeXperts (Liu et al., 2021) trains experts and anti-experts by finetuning small language models, and using these experts to guide generation. Reward-Augmented Decoding (RAD) (Deng and Raffel, 2023) trains a reward model to score generations and adjust logit probabilities to promote high-reward tokens.

heading: 7 Conclusion
text: 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.

heading: 8 Limitations
text: While DGLM demonstrates strong capabilities for guided text generation, we acknowledge important limitations. First, like any system that controls text attributes, it risks potential misuse to steer language in harmful directions. Researchers and practitioners should carefully evaluate generation systems to mitigate these risks.
In addition, DGLM currently has slower inference speed than some plug-and-play baselines when generating short texts (<32 tokens). We expect advances in accelerating diffusion models and distilling diffusion steps will help address this limitation in future work.
More broadly, while DGLM outperforms recent methods, there is still substantial room for improvement in controllable text generation. The framework currently utilizes simple linear classifiers that may not robustly capture complex attributes. Extending DGLM to complex attributes may require more complex classifiers. We hope our work sparks further research towards reliable and beneficial guided language models.

heading: Acknowledgements
text: This research is supported by grants from the National Science Foundation NSF (IIS-2107161, and IIS-1724282, HDR-2118310), the Cornell Center for Materials Research with funding from the NSF
MRSEC program (DMR-1719875), DARPA, arXiv, LinkedIn, and the New York Presbyterian Hospital.

heading: A Additional Figures
text: We present toxicity mitigation results with the Toxic Rate metric in Figure 7. We present plug-and-play results with positive sentiment guidance in Figure 8. We present the sentiment guidance results with the DistilBERT classifier used in prior work in Figure 9 and Figure 10.

heading: B Numerical Results
text: We provide the numerical results for our toxicity mitigation and sentiment control experiments in Table 5 and Table 6.
For the baseline methods, we observed a handful of extremely high perplexity generations (e.g. >1e4) that significantly increase the average perplexity. Prior work typically filters out these instances when computing the average perplexity6. We did not observe any such high perplexity continuations for our method. We therefore do not perform this filtering for our method.
C Implementation Details
We train all of the models in this work on a single NVidia A6000 GPU.
Transformer Implementation. We use different configurations of the same transformer architecture for the prompt generator and the diffusion network. We utilize a pre-normalization transformer (Vaswani et al., 2017; Xiong et al., 2020) with RMSNorm (Zhang and Sennrich, 2019) and SwiGLU activations (Shazeer, 2020). We condition the transformer on the level of noise by mapping αt to a sinusoidal positional embedding (Vaswani et al., 2017) and pass it through an MLP with a single hidden layer to obtain a time embedding. We apply adaptive RMSNorm conditioned on this time embedding before the feedforward layers and attention layer (Peebles and Xie, 2022). We utilize query-key RMSNorm (Dehghani et al., 2023) for the self-attention mechanisms because it has been shown to improve stability.
Diffusion Network. We employ the vparameterization and minimize:
Lv(x) = Et,x,ϵ[w(λt) · ∥v̂θ(zt;λ)− vt∥22].
To set the weighting function, we followed the advice of (Karras et al., 2022) and parameterized it with a log-normal distribution based on the noise
6See here for an example.
levels where the model was best able to minimize the loss. This led us to set w(λt) = N (λt; 0, 2.4). Consistent with past work (Balaji et al., 2022), we observed that increasing weights at high noise levels improved the alignment of generations with the conditioning information. For our final weighting function, we therefore used a fat-tailed Cauchy distribution for the left half of the distribution and a normal distribution for the right half. This gives us
w(λt) =
{ 1 Zc
Cauchy(λt; 0, 2.4) if λt < 0 1 Zn N (λt; 0, 2.4) if λt ≥ 0
where Zc and Zn are normalization constants such that the density of each distribution at 0 is re-scaled to 1. For training, we utilize the adaptive noise scheduler introduced by Kingma and Gao (2023) to reduce the variance of the loss estimate.
Sampling Configuration. We use the stochastic DDPM sampler with 50 sampling steps with the cosine noise schedule (Nichol and Dhariwal, 2021). We follow Hoogeboom et al. (2023) and set the variance for the DDPM sampler to a logscale interpolation between the upper and lower bounds of the variance from Ho et al. (2020): σ2 = exp(v log(σ2max) + (1− v) log(σ2min)) with v = 0.2 . We did not explore this choice in detail and further exploration of sampling configurations would likely improve performance.
Logistic Regression Classifiers. We train logistic regression models with sci-kit learn. We utilize the default L-BFGS solver with L2 regularization of 1e-3. We use balanced class weights for the toxicity classifier due to the class imbalance in the toxicity dataset. Our toxicity and sentiment classifiers achieve an Area Under the Receiver Operating Curve of 83.7 and 95.7, respectively.

heading: D Additional Composition Results
text: As specified in the main paper, DGLM is naturally suited for simultaneously controlling multiple attributes. Table D presents additional qualitative results for compositional control. From the table, we see that the instances satisfies both control attributes.

heading: Future Work
text: We expect advances in accelerating diffusion models and distilling diffusion steps will help address this limitation in future work. More broadly, while DGLM outperforms recent methods, there is still substantial room for improvement in controllable text generation. The framework currently utilizes simple linear classifiers that may not robustly capture complex attributes. Extending DGLM to complex attributes may require more complex classifiers. We hope our work sparks further research towards reliable and beneficial guided language models