D iffusion BERT : Improving Generative Masked Language Models with Diffusion Models (2023)

Abstract: We present DiffusionBERT, a new generative masked language model based on discrete dif- fusion models. Diffusion models and many pre- trained language models have a shared training objective, i.e., denoising, making it possible to combine the two powerful models and enjoy the best of both worlds. On the one hand, dif- fusion models offer a promising training strat- egy that helps improve the generation quality. On the other hand, pre-trained denoising lan- guage models (e.g., BERT) can be used as a good initialization that accelerates convergence. We explore training BERT to learn the reverse process of a discrete diffusion process with an absorbing state and elucidate several designs to improve it. First, we propose a new noise schedule for the forward diffusion process that controls the degree of noise added at each step based on the information of each token. Sec- ond, we investigate several designs of incorpo- rating the time step into BERT. Experiments on unconditional text generation demonstrate that DiffusionBERT achieves significant improve- ment over existing diffusion models for text (e.g., D3PM and Diffusion-LM) and previous generative masked language models in terms of perplexity and BLEU score. Promising re- sults in conditional generation tasks show that DiffusionBERT can generate texts of compa- rable quality and more diverse than a series of established baselines.

## introduction diffusion models @xcite @xcite have recently emerged as a new class of state-of-the-art generative models, achieving high-quality synthesis results on image data @xcite @xcite . though these models captured widespread attention from * equal contribution. † corresponding author. this work, to our knowledge, is the first attempt to combine diffusion models with plms. such a combination is built upon a shared training ob-jective between diffusion models and plms, i.e., denoising. diffusion models consist of a forward process (data to noise) and a reverse process (noise to data). in the forward process, a small amount of noise is gradually added to the data. then, a neural network (p θ in figure 1 ) is employed to learn the reverse process step by step, i.e., learn to denoise. such a denoising neural network is naturally related to a wide class of plms that are pre-trained with denoising objectives such as bert @xcite and bart @xcite . hence, pre-trained denoising language models can serve as a good starting point to learn the reverse diffusion process. on the other hand, diffusion models also offer a promising training strategy for generative plms. in contrast to commonly used generative plms (e.g., @xcite ) that rely on an autoregressive factorization of the joint probability, diffusion models provide another way of factorization along the dimension of time and therefore allow the generative model to be not necessarily autoregressive. thus, diffusion models can be combined with a variety of plms that may not be pre-trained for generation. in the discrete domain, the forward diffusion process can be implemented by a chain of transition matrices that gradually corrupt the clean text. as shown in figure 1 , the clean text "hello world !" is gradually corrupted into "[mask] [mask] [mask]" during the diffusion process. in this work, we explore using pre-trained denoising language models (e.g., bert) to learn the reverse diffusion process and demonstrate their advantages in accelerating convergence and improving generation quality. further, we propose a new noise schedule of the forward process based on the principle of distributing the corrupted information uniformly across the forward process. the noise schedule, called spindle schedule, generates noise for x t conditioned not only on x t-1 but also on x 0 , making the forward process non-markovian without changing the original training objective. note that the denoising model takes as input x t and time step t to predict x t-1 , where t is unseen during the pre-training of language models so we investigate several ways of incorporating the time step into plms. as a result, we find that the best result is achieved by throwing away the time information, which we call time-agnostic decoding (tad). experimental results on unconditional text generation demonstrate the benefit of combining dif-fusion models with plms: the proposed dif-fusionbert significantly improves the generation quality over existing diffusion models for text generation (e.g., d3pm @xcite and diffusion-lm @xcite ) and previous generative masked language models (e.g., bert-mouth (wang and cho, 2019)). diffusion-bert also matches several strong baselines in conditional generation tasks and shows superior generation diversity. the effectiveness of the proposed spindle schedule and time-agnostic decoding is confirmed by ablation studies. in a nutshell, dif-fusionbert enjoys the best of both worlds. ## diffusionbert in contrast to recently proposed diffusion models for text, e.g., diffusion-lm @xcite and diffuseq @xcite , which are based on continuous diffusion models, we instead explore discrete diffusion models to integrate plms as the backbone. we first introduce a specific instance of discrete diffusion models @xcite , which considers a transition matrix with an absorbing state for the sake of using plms ( § 3.1). secondly, we introduce a new noise schedule of the forward diffusion process, called spindle schedule, which is based on the principle of distributing the corrupted information uniformly across the forward process ( § 3.2). then, we investigate several alternatives of incorporating the time step into plms for predicting x t-1 given x t and t ( § 3.3). finally, we explore training diffusionbert for conditional generation with prompts ( § 3.4). ## conclusion this work aims to approach the problem of text generation with non-autoregressive models. to achieve this, we combine pre-trained denoising language models with absorbing-state discrete diffusion models. the training procedure of our method includes two main deviations from current discrete diffusion models, i.e., a new family of time step designs and the spindle noise schedule. the spindle noise assigns a schedule for each token according to its frequency in the training corpus. experimental results of unconditional generation demonstrate the success of diffusionbert in terms of genera-tion quality and diversity. in constrained settings, diffusionbert surpasses 7 strong baselines in generation variety by a large margin and its generation quality matches state-of-the-art methods.