Source-linked AI summary
SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Control
Xiaochuang Han, Sachin Kumar, Yulia Tsvetkov
TL;DR
General-purpose diffusion language models have lagged autoregressive models and face constraints in flexible decoding and modular control. SSD-LM combines semi-autoregressive block generation with simplex-based diffusion and off-the-shelf classifier guidance. It matches or outperforms strong GPT-2 baselines on unconstrained generation and surpasses competitive controlled-generation baselines, while decoding more slowly than autoregressive models.
Problem
Diffusion language models have underperformed autoregressive models on general-purpose text and face challenges with flexible output length and modular control.
Method
SSD-LM generates token blocks semi-autoregressively and denoises vocabulary simplexes, enabling variable-length decoding and modular guidance from shared-tokenizer classifiers.
Results
SSD-LM matches or outperforms strong GPT-2 models on unconstrained generation and surpasses competitive baselines on controlled text generation.
Takeaways & Limitations
SSD-LM provides flexible and modular diffusion-based text generation while outperforming diffusion baselines and competing controlled-generation methods.
Takeaways & Limitations
SSD-LM decodes considerably more slowly than autoregressive models; generating a 25-token block from 50-token context takes 25 seconds in one unbatched setting.
Abstract
from arXiv · showhide
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.
1 Introduction
Diffusion models remain less effective than autoregressive language models for general-purpose text, while facing challenges in flexible-length generation and modular control. SSD-LM addresses these challenges with semi-autoregressive blocks and simplex-based diffusion, achieving strong benchmark results.
- Diffusion models for text still underperform autoregressive language models despite success in continuous-valued domains.
- Semi-autoregressive generation produces left-to-right token blocks with bidirectional context, supporting variable-length sequences and refinement within blocks.
- Simplex-based diffusion represents discrete text as vocabulary distributions and enables modular guidance from off-the-shelf classifiers sharing the tokenizer.
- SSD-LM matches or outperforms strong GPT-2 models on unconstrained generation quality and diversity while surpassing competitive controlled-generation baselines.
2 Background
Diffusion models iteratively denoise continuous data, whereas autoregressive language models predict discrete tokens sequentially. SSD-LM adapts diffusion to text by denoising vocabulary simplexes and generating token blocks.
- Diffusion model: Canonical diffusion training adds Gaussian noise across timesteps, then learns to reverse the process by predicting the added noise.
- Diffusion model: Diffusion inference starts from Gaussian noise and iteratively reconstructs an output by stepping backward through the timesteps.
- Diffusion model: Text requires an adaptation because tokens are discrete, so SSD-LM represents them as continuous vocabulary simplexes before applying diffusion.
- Autoregressive LM: Autoregressive language models optimize sequence likelihood by predicting and sampling each next token conditioned on preceding context.
- SSD-LM: Figure 1 depicts token order horizontally and diffusion timesteps vertically, distinguishing observable variables, discrete vocabulary, continuous logits, and model inputs.
- SSD-LM: SSD-LM decodes blocks of B tokens conditioned on context, combining semi-autoregressive generation with diffusion-based refinement.
3 SSD-LM
SSD-LM combines semi-autoregressive block generation with simplex-based diffusion over vocabulary representations. Its decoding repeatedly refines token blocks, supports classifier guidance, and permits flexible decoding choices.
- Model design: SSD-LM generates a block of B tokens from noisy representations conditioned on the preceding context w<c.The model adapts autoregressive context conditioning and diffusion-based block refinement.
- Simplex representation: Tokens are represented as almost-one-hot probability simplexes over the vocabulary, without learnable parameters in the mapping.Forward diffusion adds time-dependent Gaussian noise to the token logits, producing fully noisy simplexes at the final step.
- Training objective: The training objective estimates conditional likelihood for original tokens from noisy vocabulary simplexes rather than using an L2 noise-prediction loss.The objective conditions on context and randomly sampled block positions and diffusion timesteps; pilot results found L2 distance ineffective, potentially because the simplex representation is intrinsically skewed.
- Decoding: During decoding, reverse diffusion predicts token logits, projects them back toward almost-one-hot data representations, and selects the final block with argmax.Projection can be greedy, top-p sampling, or multi-hot; generated blocks are appended to the context before decoding the next block.
- Decoding: Unlike standard DDPM decoding, SSD-LM uses a novel decoding algorithm because vanilla DDPM decoding failed to produce sensible simplex-based outputs in preliminary experiments.Generation can repeat block decoding until the desired length or an end-of-sequence token is reached.
- Modular control: Classifier guidance uses frozen standalone attribute models directly when they share the language model’s tokenizer, enabling modular controlled generation without classifier retraining.A control-weight hyperparameter balances the guidance, and classifiers are modified only to accept simplex inputs.
4 Experiments
Experiments evaluate SSD-LM against autoregressive and diffusion-based language models on unconstrained generation, ROCStories, and controlled sentiment generation. SSD-LM shows strong quality, diversity, controllability, and modularity, though perplexity involves trade-offs.
- 4.1 SSD-LM pretraining setup: SSD-LM uses RoBERTa-large with BPE tokenization and is pretrained on 9B OpenWebText tokens for comparison with GPT-2.The model is trained for 100K steps on 32 Nvidia V100 GPUs, taking about six days.
- 4.1 SSD-LM pretraining setup: SSD-LM’s average pretraining NLL reaches 3.87, while canonical training-time perplexity is not directly compatible with diffusion LMs.The loss is a per-token NLL that depends on the selected noise schedule.
- 4.2 Unconstrained text generation: SSD-LM generates 25-token decoding blocks with configurable logits projection and sampling strategies, extending blockwise generation to longer outputs.For 50- or 100-token generations, each new block is fed as context for the next block.
- 4.2 Unconstrained text generation: SSD-LM outperforms larger GPT-2 models on MAUVE and diversity metrics, including Dist-n, while achieving lower repetition rates.Perplexity results are mixed, with a trade-off between MAUVE and perplexity across settings.
- 4.2 Unconstrained text generation: On ROCStories, SSD-LM achieves substantially higher MAUVE and lower perplexity than Diffusion-LM on held-out 50-token continuations.The evaluation uses 2,700 held-out sequences with 50-token prompts and 50-token outputs.
- 4.3 Controlled text generation: For controlled sentiment generation, SSD-LM provides strong controllability and modularity, outperforming M&M LM while balancing control, fluency, and diversity against CC and CM methods.Evaluation considers external classifier accuracy, perplexity, and distinctiveness across output lengths.
5 Related work
Related work contrasts continuous-domain diffusion success with the challenges of applying diffusion to discrete text and controlling attributes. Prior approaches trade off flexibility, generality, or modularity in different ways.
- Diffusion models: Diffusion models perform strongly in continuous domains, but Gaussian-noise formulations are not straightforward to apply to discrete text.Prior text approaches use approximations such as embeddings, character-level, or byte-level representations.
- Diffusion models: Earlier diffusion language models support blockwise length flexibility but were trained on specialized datasets and not evaluated against general-purpose autoregressive models for unconstrained generation.Some embedding-based methods also require customized attribute classifiers for post-hoc control.
- Language models: Autoregressive language models predict tokens left-to-right and have achieved strong benchmark capabilities, but their decoding paradigm limits flexible attribute control.Non-autoregressive methods provide broader context but introduce other design challenges.
- Controllable text generation: Control-code approaches require retraining models for new controls, motivating decoding-time methods that alter output distributions without modifying pretrained autoregressive models.The related work includes classifier- and objective-based approaches to controlled generation.
6 Conclusion
SSD-LM combines semi-autoregressive block decoding with simplex diffusion over the output vocabulary. It matches or exceeds strong GPT-2 baselines, surpasses diffusion baselines, and supports modular controlled generation.
- Conclusion: SSD-LM is a semi-autoregressive diffusion language model trained to denoise corrupted simplexes over the output vocabulary.It generates blocks of text, providing flexible output length and refinement within blocks.
- Conclusion: On unconstrained generation, SSD-LM performs on par with or outperforms strong, larger GPT-2 models in quality and diversity while vastly outperforming Diffusion-LM.The conclusion summarizes results across the paper’s unconstrained generation evaluations.
- Conclusion: On controlled generation, SSD-LM surpasses baselines while retaining an easy-to-use modular design based on off-the-shelf attribute classifiers.The paper frames flexible and modular diffusion-based language generation as a direction for future research.
Limitations
The limitations concern training efficiency, decoding speed and scheduling, the scale and diversity of controls, and the primarily English pretraining setting.
- Sample efficiency: SSD-LM has lower sample efficiency than autoregressive LMs because each training example computes loss on only B tokens rather than all L tokens.The authors suggest architectures dedicated to semi-autoregressive diffusion as a possible improvement.
- Decoding speed: SSD-LM decoding is considerably slower than autoregressive decoding because each block is refined over multiple iterations.With a 50-token context, generating one 25-token block takes 25 seconds for an unbatched instance at Tdecode=1000.
- Decoding block size: The decoding block size remains constant across iterations, leaving the schedule and block lengths as predefined rather than learned.The authors propose exploring flexible schedules and learning Bdecode and Btrain.
- Scope: The evaluation covers limited control combinations and primarily English pretraining data, leaving larger-scale, multilingual, and broader control settings for future work.The authors also identify alternative continuous representations and noise processes as future directions.
Ethics statement
SSD-LM remains exposed to harms documented for autoregressive language models, including bias, toxicity, factual inconsistency, disinformation, and privacy risks. Its controllability may also be maliciously misused.
- SSD-LM is subject to risks including biased or toxic generation, factual inconsistency, disinformation, and compromised user privacy.
- Controllability could potentially alleviate some harms, but it can also be misused to generate biased, toxic, or non-factual content.
A A contrastive interpretation of the training loss
The SSD-LM training objective can be interpreted contrastively: it favors true data likelihood while penalizing noisy-data likelihood across noise scales.
- The objective maximizes estimated likelihood of true data while penalizing estimated likelihood of noisy data.This interpretation differs from the exact DDPM objective supported by a variational bound.
- The contrastive interpretation compares true-data and noisy-data likelihoods over a broad range of noise scales.
B Connection between our decoding algorithm and the DDPM decoding
SSD-LM’s decoding procedure parallels DDPM decoding by predicting an initial data representation and then applying a forward diffusion step, while sampling noise to encourage exploration.
- The decoding transition is derived by rearranging the DDPM decoding transition and comparing its form with the forward diffusion equation.
- DDPM decoding can be viewed as predicting an imaginary x_0 and applying a compensating forward diffusion step with deterministic predicted noise.
- SSD-LM uses a similar predict-then-diffuse structure but samples noise z instead of using deterministic predicted noise.The sampled noise is intended to encourage exploration.
- For most timesteps, 1−ᾱ_t > 0.98, with outliers as t approaches 0 or T.
C Detailed setup of the comparison with Diffusion-LM (Li et al., 2022)
The Diffusion-LM comparison uses matched experimental setups with specified data, optimization, tokenization, and decoding configurations. A BERT-initialized Diffusion-LM produced severely degraded results relative to both baselines.
- 50K ROCStories training sequences contain 100 tokens each after block concatenation, while both models use BERT-base-uncased tokenization.
- SSD-LM uses batch size 512, learning rate 1e-4, 20K steps, block-wise generation with m=2 iterations, Tdecode=2500, and top-p=0.5.
- Diffusion-LM uses batch size 64, learning rate 1e-4, 400K steps, and is evaluated on an infilling task with Li et al.’s decoding hyperparameters.
- A BERT-initialized Diffusion-LM obtained MAUVE 0.4 out of 100 and PPL 73157, much worse than the default Diffusion-LM and SSD-LM.All checkpoints showed the same degenerated result, which the authors conjecture reflects incompatibility with pretrained non-diffusion weights.
D Additional results
Additional analyses show how decoding projections and control settings affect SSD-LM, while tables and trajectories provide evaluations and qualitative generation evidence.
- Logits projection: Greedy projection lowers perplexity but increases repetition, making it undesirable for unconstrained generation.Reducing top-p toward 0 produces low perplexity alongside a high repetition rate.
- Logits projection: Multi-hot projection performs worse overall than sampling projection, favoring intermediate states committed to single tokens.The authors suggest multi-hot projection may become viable with multi-hot logits mapping for input tokens.
- Controlled generation: Larger control weight improves external classifier accuracy but also increases perplexity, with reasonable behavior for top-p values of 0.2 and 0.5.Figure 5 varies control weight λ and top-p across generation lengths, decoding strategies, and decoding timesteps.
- Evaluation results: Tables 4 and 5 evaluate unconstrained generation at lengths 25 and 100 using PPL from GPT-Neo-1.3B, with configurations ranked by MAUVE.GPT-2 results are averaged across five seeds, while SSD-LM reports its top three configurations.
- Evaluation results: Tables 6 and 7 report controlled generation at lengths 12 and 20 using external and, when available, internal classifier accuracy plus perplexity.Perplexity is computed with GPT2-xl, and baselines differ in classifier and language-model customization.
- Qualitative analysis: Tables 8 and 9 provide qualitative generations and intermediate decoding states, including unconstrained and sentiment-controlled examples.Table 9 tracks intermediate states as t decreases for a 25-token block and top-p sampling of 0.99.