[Paper Review] TESS: Text-to-Text Self-Conditioned Simplex Diffusion

Rabeeh Karimi Mahabadi, Hamish Ivison+5 more|arXiv (Cornell University)|May 15, 2023

Topic Modeling4 citations

TL;DR

TESS proposes a fully non-autoregressive text diffusion model that performs diffusion directly on the logit simplex space, using a novel self-conditioning mechanism to improve generation quality. It outperforms state-of-the-art non-autoregressive and diffusion-based models on summarization, paraphrase generation, text simplification, and question generation, achieving performance competitive with strong autoregressive models like BART.

ABSTRACT

Diffusion models have emerged as a powerful paradigm for generation, obtaining strong performance in various continuous domains. However, applying continuous diffusion models to natural language remains challenging due to its discrete nature and the need for a large number of diffusion steps to generate text, making diffusion-based generation expensive. In this work, we propose Text-to-text Self-conditioned Simplex Diffusion (TESS), a text diffusion model that is fully non-autoregressive, employs a new form of self-conditioning, and applies the diffusion process on the logit simplex space rather than the learned embedding space. Through extensive experiments on natural language understanding and generation tasks including summarization, text simplification, paraphrase generation, and question generation, we demonstrate that TESS outperforms state-of-the-art non-autoregressive models, requires fewer diffusion steps with minimal drop in performance, and is competitive with pretrained autoregressive sequence-to-sequence models. We publicly release our codebase at https://github.com/allenai/tess-diffusion.

Motivation & Objective

• To address the challenges of applying diffusion models to discrete natural language by avoiding autoregressive generation and embedding space dependencies.
• To improve text generation quality in non-autoregressive settings through a novel self-conditioning mechanism on the logit simplex space.
• To demonstrate that direct diffusion on the probability simplex enables flexible, global sequence generation without block-wise constraints.
• To evaluate the model across diverse NLP tasks, including both generation and understanding, showing competitive performance against strong autoregressive baselines.
• To release trained models and code to advance open research in diffusion-based text generation.

Proposed method

• TESS performs diffusion directly on the vocabulary logit space (simplex space), avoiding learned embeddings and enabling end-to-end training on discrete tokens.
• It employs a novel form of self-conditioning that leverages the geometric semantics of the simplex space to improve denoising accuracy and generation quality.
• The model uses a standard diffusion training objective: denoising the noise added to the logit vectors via a U-Net architecture, with noise schedule and reverse process modeling.
• Self-conditioning is applied during inference by conditioning the reverse denoising process on the current state of the sequence, enhancing coherence and fluency.
• The method is trained end-to-end from a RoBERTa checkpoint, fine-tuning on downstream tasks without additional pretraining.
• Inference uses 1000 sampling steps for generation and 10 for classification, with a simplex scale hyperparameter $k=5$.

Experimental results

Research questions

• RQ1 Can a fully non-autoregressive diffusion model achieve competitive performance on natural language generation tasks without relying on autoregressive generation?
• RQ2 Does self-conditioning on the logit simplex space improve generation quality compared to standard self-conditioning or no conditioning?
• RQ3 Can direct diffusion on the probability simplex outperform diffusion in embedding space or other latent spaces in terms of both quality and flexibility?
• RQ4 How does TESS compare to strong autoregressive models like BART on diverse NLG and NLU tasks?
• RQ5 Does the proposed simplex-based diffusion framework generalize across multiple NLP tasks, including summarization, paraphrase generation, and text simplification?

Key findings

• TESS outperforms all non-autoregressive and diffusion-based baselines on paraphrase generation, achieving higher BLEU, BERTScore, and ROUGE-L scores.
• On summarization, TESS achieves ROUGE-L F1 scores competitive with BART, despite not using the BART pretraining objective.
• In text simplification, TESS achieves SARI scores significantly higher than prior diffusion models and competitive with strong autoregressive models.
• For question generation, TESS surpasses all non-autoregressive and diffusion baselines in both quality and diversity metrics.
• On the GLUE benchmark, TESS performs comparably to strong masked language models, demonstrating strong zero-shot and few-shot capabilities on natural language understanding tasks.
• The proposed simplex-based self-conditioning method leads to a substantial improvement in generation quality, as validated across all evaluated tasks.