A Cheaper and Better Diffusion Language Model with Soft-Masked Noise (2023)

Abstract: Diffusion models that are based on iterative denoising have been recently proposed and leveraged in various generation tasks like image generation. Whereas, as a way inherently built for continuous data, existing diffusion models still have some limitations in modeling discrete data, e.g., languages. For example, the generally used Gaussian noise can not handle the discrete corruption well, and the objectives in continuous spaces fail to be stable for textual data in the diffusion process especially when the dimension is high. To alleviate these issues, we introduce a novel diffusion model for language modeling, Masked-Diffuse LM, with lower training cost and better performances, inspired by linguistic features in languages. Specifically, we design a linguistic-informed forward process which adds corruptions to the text through strategically soft-masking to better noise the textual data. Also, we directly predict the categorical distribution with cross-entropy loss function in every diffusion step to connect the continuous space and discrete space in a more efficient and straightforward way. Through experiments on 5 controlled generation tasks, we demonstrate that our Masked-Diffuse LM can achieve better generation quality than the state-of-the-art diffusion models with better efficiency.

## introduction we present a novel diffusion method for modeling languages, masked-diffusion lm (language model), which uses strategic soft-masking informed by linguistic features to corrupt both the discrete and continuous space, and then iteratively denoise them back by predicting the categorical distribution. specifically, a strategic soft-masking process is designed that gradually adds perturbation to the input text in an order from harder or more informative words to simpler or less informative words through soft-masking. as a result, the models are encouraged to recover and generate the text following an easy-first-generation nature @xcite to improve the generation structure and quality with more flexibility. also, during the diffusion process, we directly predict the discrete token with cross-entropy loss that maps the continuous space to discrete textual space to stabilize the intermediate diffusion steps. through our proposed masked-diffusion lm, the applicationspecific performance metrics as well as training efficiency are significantly improved over current diffusion language models based on experiments. our work is inspired by recent advances in diffusion models @xcite @xcite @xcite that are introduced as a new generative modeling approach based on iterative denoising and have achieved high-quality generations for visual and audio modalities @xcite @xcite @xcite . although these approaches have received growing attention and achieved impressive success, applying diffusion models to textual domain is still challenging and under-explored due to the discrete nature of the text (e.g., one-hot vectors) compared to continuous data like images (e.g., rgb values) @xcite . a few prior works @xcite @xcite @xcite that explore using diffusion models on textual data can be divided into two lines. the first is to extend diffusion models to discrete state spaces @xcite . the second is to perform the diffusion process and its reverse process in the continuous domain and bridge the continuous and the discrete domain through embedding and rounding @xcite , for example, diffusion-lm @xcite . despite the improvements, most previous works fail to leverage the linguistic features (e.g., words in sentences are with different importance) to noise the input textual data and recover it back in a more suitable way. besides, they usually neglect or fail to adapt large pre-trained language models (plms) @xcite @xcite @xcite @xcite @xcite , which is an unmissable treasure in the nlp community: their adopted k-nearest-neighbor rounding technique that maps continuous space to discrete space cannot handle high-dimensional data in a stable and efficient way @xcite . as a result, a corruption process tailored for languages and the objective that allows efficient and straightforward discrete and continuous space transformation is in great need. our masked-diffusion lm realizes this extension. to demonstrate the effectiveness of our introduced masked-diffusion lm, we perform experiments on e2e dataset @xcite and 5 controllable generation tasks @xcite including semantic content, parts-of-speech, syntax tree, syntax spans, and length. we observe that our masked-diffusion lm can (i) achieve the stateof-the-art performances compared to recent baseline models, and (ii) allow more efficient training and inference compared to previous diffusion-lm. to summarize, our contributions are: (1)we introduce a strategic masking noise strategy guided by linguistic features to corrupt the textual data in diffusion models for modeling languages. (2) we use linear layers and cross-entropy objectives to bridge the continuous and discrete spaces in the diffusion process for efficiency and stability. (3) we conduct experiments on different controllable generation tasks to demonstrate the effectiveness of our proposed methods compared to previous diffusion language models. non-autoregressive text generation most language models @xcite ) and text generation models @xcite @xcite follow a left-to-right autoregressive manner. however, the fixed generation order prevents the models' flexibility in editing former text based on later generation results, especially for global controllable generation settings. to overcome the limitations, non-autoregressive text modeling has been proposed @xcite @xcite @xcite through masked language models @xcite , iterative sequence alignment @xcite , insertion and deletion @xcite , or unrolling the generation path @xcite . our masked-diffusion lm achieves the non-autoregressive generation through gradually recovering the intermediate latent variables in a planned sequence from the forward process. plug-and-play controllable generation our work is also closely related to the line of research about plug-and-play controllable generation methods @xcite @xcite , which modify the outputs based on extra guidance such as classifiers without changing or fine-tuning the pretrained language models. @xcite used gradients to edit the autoregressive language model's hidden representations to fulfill the control guidance. @xcite proposed to reweight the predicted token from the language models while @xcite further fine-tuned a smaller lm to reweight the token predictions. in this work, we apply the gradient-based plug-and-play approach to our masked-diffusion lm for controllable generation by making classifier-guided gradient updates to the intermediate latent variables during the diffusion. ## method: the masked-diffusion lm in this section, we describe our introduced masked-diffusion lm. the overall diagram is shown in figure 1 and algorithm 1,2. different from the recent diffusion models for languages, e.g., diffusion-lm @xcite , which are based on continuous diffusion models, we propose to make corruptions in both discrete and continuous space to help modeling the textual data. specifically, we formulate a novel corruption process as an alternative to gaussian diffusion (in section 3.2) and we directly map continuous vectors to discrete inputs in every diffusion step with cross-entropy objectives (in section 3.3). moreover, our approach could easily integrate pre-trained language models (in section 3.4). ## controllable text generation with masked-diffusion lm in this section, we illustrate how we apply our masked-diffusion lm to fulfill controllable text generation. inspired by recent plug-and-play methods @xcite @xcite , we conduct controls c from external modules (e.g., classifiers) directly on the latent variables x t in every intermediate step @xmath0, t ] in our masked-diffusion lm: we follow the conditional independence assumption @xcite @xcite and decompose the above joint probability into a sequence of control task at every time step t: for the decoding strategy, following li et al. (2022), the minimum bayes risk (mbr) decoding @xcite ) is used to aggregate and select the sample that has the lowest expected loss under the specified loss function from the masked-diffusion lm. ## conclusion in this work, we present a novel diffusion model for language, masked-diffusion lm, which corrupts the discrete text with a linguistic-informed softmasking strategy and then iteratively denoises them back by directly predicting the text. specifically, we gradually soft-mask the tokens in the sentence following an order from more informative words to less informative words in the forward process. this satisfies the flexibility for diffusion models, as well as encourages the easy-first-generation nature in the denoising process for better generation quality. also, we directly predict the discrete token during the diffusion process with the cross-entropy loss to stabilize the intermediate diffusion steps and make our approach orthogonal to large pre-trained language models. experiments on e2e dataset and five controllable generation tasks including semantic content, parts-of-speech, syntax tree, syntax spans, and length show that our masked-diffusion lm can (i) achieve the state-of-the-art performances compared to recent baseline models and (ii) allow more efficient training and inference compared to the previous diffusion-lm.