Empowering Diffusion Models on the Embedding Space for Text Generation (2024)

Abstract: Diffusion models have achieved state-of-the-art synthesis quality on both visual and audio tasks, and recent works further adapt them to textual data by diffusing on the embedding space. In this paper, we conduct systematic studies of the optimization challenges encountered with both the embedding space and the denoising model, which have not been carefully explored. Firstly, the data distribution is learnable for embeddings, which may lead to the collapse of the embedding space and unstable training. To alleviate this problem, we propose a new objective called the anchor loss which is more efficient than previous methods. Secondly, we find the noise levels of conventional schedules are insufficient for training a desirable denoising model while introducing varying degrees of degeneration in consequence. To address this challenge, we propose a novel framework called noise rescaling. Based on the above analysis, we propose Difformer, an embedding diffusion model based on Transformer. Experiments on varieties of seminal text generation tasks show the effectiveness of the proposed methods and the superiority of Difformer over previous state-of-the-art embedding diffusion baselines.

## introduction a wave of diffusion models @xcite @xcite is sweeping the generation tasks (e.g., image and audio synthesis) recently, showing their great capacity for high-quality data generation. diffusion models are a family of iterative generative models, which are trained to recover corrupted data and then generate data by gradually refining samples from the pure noise. this procedure enables the model to make subtle refinements of output samples in a * equal contribution. the booming achievements in vision and audio domains inspire researchers to delve into the realm of text generation. diffusion models introduce a novel noising paradigm and a training objective other than token prediction, establishing an alternative form of language models, which exhibits the potential to foster an enhanced comprehension of language modeling. from a higher perspective, this investigation generalizes the diffusion model across modalities, and further contributes to a unified multimodal framework @xcite . nonetheless, the exploration is still at an initial stage. recent works @xcite @xcite basically convert the discrete tokens to embeddings and then utilize continuous diffusion models to generate them, which can be termed embedding diffusion models. these preliminary attempts follow the original model to deal with the embeddings, with little consideration of the unique properties and the optimization challenges of the embedding space and the denoising model. in this paper, we explore the embedding diffusion model from two perspectives separately, i.e., the embedding space and the denoising model, based on which we conduct a thorough study respectively. firstly, for diffusion models on image and audio generation, the ground truth data is stationary during training. in contrast, it is learnable for the textual data (i.e., embeddings), which may cause the collapse of the embedding space and introduce instability to the training of the model. to avoid the collapse caused by dynamically shifting embedding parameters, we propose an anchor loss function to attain well-distributed embeddings and stabilize the training process. the detailed analysis is presented in section 3.1. secondly, in section 3.2, we find that in the high dimensional embedding space, the insufficient noise results in a simple denoising task, which causes the degeneration of the model. to tackle this challenge, we propose a novel framework named noise rescaling, which is orthogonal to the choice of the noise schedule and applicable to any existing schedules. specifically, we define an index termed degeneration score as a measurement of the degree of degeneration. guided by the degeneration score, we can apply a noise rescaling procedure to prevent the model from degenerating. based on the above discussion, we propose an integrated framework of difformer, a denoising diffusion transformer model. we conduct experiments on a variety of important text generation tasks including machine translation, text summarization, paraphrasing, text simplification, and question generation. on these benchmark datasets, difformer outperforms diffusion-based and iteration-based non-autoregressive baselines and achieves state-ofthe-art performance among embedding diffusion models. further experiments demonstrate the superiority of difformer over baselines including llms in quality, diversity, and efficiency, emphasizing the potential of diffusion models for text generation in the era of llms. ## methodology this section elucidates the challenges inherent in optimizing embedding diffusion models and presents our corresponding solutions. we start with an introduction to the model architecture. the model architecture is based on transformer @xcite , which consists of an encoder and a decoder. the decoder, as the main stem component, is considered as two separate parts in this paper, namely the embed- and the denoising model f θ (•), which denotes the stacked decoder layers. notably, this paper define z 0 = e ϕ (y). the encoder provides the representation @xmath0 ## experiments 4.1 experimental setup to evaluate the proposed difformer model, we conduct experiments on five conditional text generation tasks including machine translation, text summarization, paraphrasing, text simplification, and question generation. datasets for machine translation, mainly following previous works @xcite @xcite , three benchmark datasets wmt14 en-de @xcite , wmt16 en-ro @xcite and iwslt14 de-en @xcite are inclued. for text summarization, experiments are conducted on gigaword @xcite . in addition, following previous non-autoregressive text generation works @xcite @xcite , for machine translation and text summarization tasks, we adopt sequencelevel knowledge distillation @xcite on the original training set to alleviate the multimodality problem. for paraphrasing, text simplification, and question generation tasks, we mainly follow @xcite to conduct experiments on quora question pairs (qqp) foot_0 , wiki-auto @xcite and quasar-t (qt) @xcite respectively. the data split of the above datasets can be found in appendix e. ## conclusion in this paper, we conduct a thorough study of the challenges when optimizing an embedding diffusion model on discrete textual data, and propose the corresponding solutions. firstly, to tackle the challenge of embedding collapse and instability in training caused by the dynamic nature of embeddings, we introduce an anchor loss to regularize the embeddings and stabilize the training simultaneously. secondly, we derive a novel noise rescaling framework based on theoretical analysis, which notably alleviates the degeneration of the denoising model resulting from inadequate noise. finally, integrated with the aforementioned techniques, we present difformer, a denoising diffusion model based on transformer. difformer demonstrates superior performance on various benchmark text generation tasks, outperforming prior diffusion-based models as well as iterative non-autoregressive models.