TESS : Text-to-Text Self-Conditioned Simplex Diffusion (2024)

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.

## introduction diffusion models @xcite @xcite have achieved state-of-the-art performance in various continuous domains, such as image @xcite , audio @xcite , video @xcite , and text-to-image generation @xcite . inspired by the success of diffusion for continuous domains, recent works have adapted diffusion to discrete spaces, such as text @xcite @xcite . one line of work proposes diffusing the model latent space by adding gaussian noise to input word embeddings @xcite . another approach, ssd-lm @xcite , adds noise to the vocabulary probability simplex. direct diffusion on the probability simplex is desirable @xcite as it eliminates the need for an extra step to map diffused embeddings to actual discrete inputs or auxiliary methods such as binary encoding @xcite . despite its strong performance, however, ssd-lm has several shortcomings: a lack of self-conditioning @xcite , a lack of extensive evaluation on downstream tasks, and most notably, its restriction to generating blocks of 25 tokens, which hinders the potential benefits of full diffusion, e.g., the ability to perform arbitrary infilling, flexible generation, and a global view of the sequence. in this work, we present tess, a text-to-text diffusion model, which overcomes several limitations of prior works: restrictions on scale @xcite , dependence on pretrained embeddings @xcite , semi-autoregressive nature @xcite , and short generation length @xcite . tess closely follows @xcite by performing diffusion on the vocabulary logit space rather than the typical embedding space. unlike ssd-lm, however, tess is fully non-autoregressive and performs diffusion on the entire sequence. it also incorporates a novel form of self-conditioning, which demonstrates a competitive edge over the original self-conditioning method @xcite and dramatically improves the efficiency and quality of tess. we evaluate tess on a suite of natural language generation (nlg) tasks including summarization, text simplification, paraphrase generation, and question generation. our empirical results surpass the current state-of-the-art non-autoregressive and diffusion-based approaches and are on par with a strong pretrained encoder-decoder language model @xcite . in particular, our simplex-based self-conditioning method substantially improves generation quality. we also evaluate tess on natural language understanding (nlu) tasks from the glue benchmark @xcite and show that it performs comparably to strong masked language model baselines. our contributions can be summarized as follows. we will release our trained models and code to promote open research in the field of diffusionbased text generation. ## background we revisit continuous diffusion models @xcite , following the formulation of denoising diffusion models @xcite . training given a sample x 0 ∈ r d from a data distribution p data , a forward diffusion process q(x t |x t-1 ) is a markov chain that generates a sequence of latent variables x 1 , . . . , x t by gradually adding gaussian noise at each time step @xmath0, 2, . . . , t } with variance β @xmath1: let ϵ t ∼ n (0, i), α @xmath2, and ᾱt = t @xmath3. then sampling x t at an arbitrary time step t has the closed-form solution given a well-behaved noise schedule {β t } t @xmath4, x t follows a stationary prior distribution n (0, i). therefore, if we can approximate the reverse process q(x t-1 |x t , x 0 ) via a model p θ (x t-1 |x t ) with parameters θ, then we can sample random noise from a standard gaussian and gradually denoise it to sample from p data . in our settings, our model p θ is a transformer model foot_0 . the reverse process is thus parametrized as the model is trained by minimizing the mean squared error between the ground-truth data x 0 and its estimate xθ : foot_1 noise schedule the forward diffusion process is defined by a noise schedule. in this work, we follow the cosine schedule (nichol and dhariwal, 2021) for α t : inference in @xcite , model predictions are iteratively denoised for @xmath5, . . . , 1 starting from pure noise, following we follow the recently proposed simplex-based diffusion procedure by @xcite , which allows us to apply diffusion to text without employing auxiliary methods that map categorical data to continuous space @xcite . ## method in this section, we present tess, a simplex diffusion-based text-to-text model. building upon ssd-lm @xcite , we propose a fully non-autoregressive model with self-conditioning. continuous data representation let v denote the vocabulary space. following @xcite , we map the id of each token to be generated @xmath6, whose i-th component satisfies tess "cola sentence: the well is eating." ## conclusion we present tess, a new sequence-to-sequence diffusion model for language generation tasks that is fully non-autoregressive, works for long sequences compared to prior work, performs the diffusion process on the vocabulary logit space, and employs a new and efficient form of self-conditioning. tess outperforms strong autoregressive baselines as well as recent state-of-the-art text diffusion models on a wide variety of conditional language generation and language understanding tasks, while also being far more efficient than prior diffusion-based models. future work relies on pretraining our method combined with denoising and infilling objectives, which we hypothesize can provide further performance boosts to our text-to-text diffusion model.