Can Diffusion Model Achieve Better Performance in Text Generation ? Bridging the Gap between Training and Inference ! (2023)

Abstract: Diffusion models have been successfully adapted to text generation tasks by mapping the discrete text into the continuous space. However, there exist nonnegligible gaps between training and inference, owing to the absence of the forward process during inference. Thus, the model only predicts based on the previously generated reverse noise rather than the noise computed by the forward process. Besides, the widely-used downsampling strategy in speeding up the inference will cause the mismatch of diffusion trajectories between training and inference. To understand and mitigate the above two types of training-inference discrepancies, we launch a thorough preliminary study. Based on our observations, we propose two simple yet effective methods to bridge the gaps mentioned above, named Distance Penalty and Adaptive Decay Sampling. Extensive experiments on 6 generation tasks confirm the superiority of our methods, which can achieve \mathbf{100}\times \rightarrow \mathbf{200}\times speedup with better performance. Our code will be released at https://github.com/CODINNLG/Bridge_Gap_Diffusion .

## introduction with the prevalence of aigc (artificial intelligence generated content) in recent years, generative models @xcite have been receiving more attention. as one of the representative generative models, diffusion models @xcite have achieved great success on myriads of generation tasks with continuous data, such as image @xcite @xcite , audio generation @xcite , and molecule generation @xcite , by iteratively refining the input noise to match a data distribution. more recently, diffusion models have been successfully adapted to text generation @xcite @xcite by first leveraging an extra embedding module that maps the discrete data into the continuous space and then recovering the text from the continuous space with rounding strategy @xcite or logits projection @xcite . a typical diffusion-based text generation model contains one reverse process (from noise to data) and one forward process (from data to noise), which is shown in figure 1 . more concretely, both of the two processes can be viewed as markov chains, where the forward process gradually perturbs the data into gaussian noise while the reverse process recovers the original data step by step conditioned on the correlated noise from the forward process. the training stage involves both of the above two processes, while the inference stage only consists of the reverse process, i.e., the model predicts based on the previous noise outputted by the model itself rather than the correlated forward noise. such discrepancy between training and inference, also called exposure bias @xcite , leads to error accumulation as the denoising steps grow during the inference stage @xcite . another drawback of the diffusion model is that it requires multiple iterative denoising steps to produce the final results since the reverse process should approximate the forward process @xcite , which usually involves thousands of steps. numerous iterative reverse steps of diffusion models are inevitably time-consuming for text genera-tion. for instance, a diffusion model takes around 12 hours on one single nvidia a100 gpu to finish the inference of 10k sentences with a length of 128 while the cmlm-based non-autoregressive model @xcite only takes a few minutes 1 . to accelerate the inference speed in text generation, down sampling @xcite is leveraged @xcite @xcite , though much faster but at the cost of performance owing to the gap between the downsampled steps in inference and the full diffusion trajectory in the training stage. to explore the insights and the potential improvement of the aforementioned training-inference gaps, we conduct a preliminary study with a diffusion model @xcite on the story generation task and mainly observe that: (1) injecting the noise generated by the model itself into the training stage can improve the model performance, and (2) the uniform downsampling strategy in the inference that treats each step equally impairs the model performance, and adaptive sampling strategy should be applied for different generation stages. accordingly, we propose two simple yet effective strategies: distance penalty and adaptive decay sampling, to bridge the training-inference gaps and accelerate the inference process. experiments on 6 generation tasks of 3 different settings (directed, open-ended, and controllable) show the superiority of our methods without changing the original architecture of the diffusion model or adding more parameters. surprisingly, our methods can achieve 100× speedup with performance improvement or 200× acceleration with competitive results. ## gaps between training and inference from the above description of diffusion models, we can summarize two gaps: (1) the reverse process at time step t in inference is conditioned on the predicted noise z t+1 by the model itself while z t+1 can be obtained directly with the forward computation q(z t+1 | z 0 ) during training, and (2) the downsampled time subset t ′ in inference is inconsistent with the full diffusion trajectory t in training stage when applying the downsampling method for inference speedup. to calibrate the effects of these two types of training-inference gaps, we launch a study on the story generation task in this section. ## method we propose two simple yet effective methods: distance penalty in the post-training stage and adaptive sparse sampling in inference to bridge the gaps without introducing any architecture modification to diffusion models. thus, it can be flexibly adapted to different diffusion model variants. ## conclusion this work focuses on bridging the training and inference gaps of the diffusion model. the result of the preliminary study shows that injecting predicted noise into the model can help mitigate the gaps, and the uniform downsampling strategy for inference acceleration harms the model performance. thus, we propose two simple yet effective strategies: distance penalty and adaptive decay sampling, to mitigate the aforementioned gaps. experiments on 6 text generation tasks of 3 different settings show that the model with our methods can achieve better performance and great inference speedup.