Energy-Based Diffusion Language Models for Text Generation

[Q1] The window size and sampling size in Figure 1(a). How EDLM achieves 1.3x sampling speedup.

The window size is $0.2$ and the sampling size is $2$. Note that, the importance sampling can be conducted in parallel in GPU.

Then, just as a conceptual example, as shown in Fig (1.a), EDLM with 512 sampling steps can achieve similar results as MDLM with 1024 steps. Then, MDLM takes 1024 forward calls of the network, while EDLM takes 512*(1+0.2)≈615 calls. Since a forward pass of diffusion and energy function share a similar network and will take a similar time, the speed up will be around 1024/612=1.67. But note that this is just a conceptual example, and in practice, we get the 1.3x speedup.

[Q2] Why does adding the reranker in the early stages yield better results instead?

We think the reason is the model tends to make more errors at the start due to limited context and importance sampling will correct them. In the late stage, even if the energy can differentiate samples, the energy cannot make a difference there since all samples are bad.

Besides, we would like to mention that Figure 2 visualizes positive and negative samples from true $x_t$, e.g., $x_t$ is generated by diffusing true $x_0$ instead of by sampling. Therefore, the energy behavior is not directly the phenomenon during sampling.

[Q3] Parameterization of EDLM-NCE energy function.

It’s a separate model, though fine-tuned from the pre-trained diffusion model.

[Q4] Discrepancy about Fig 3.

Thank you for your detailed check!! We mislabeled the window sizes reversely, and it is even different from our own Fig 1. (We labeled the importance sampling ending timestep, and therefore the window should be $1-$ this value). We have updated the figure in our new version.

We hope our response could address your questions!

Thank you for your constructive feedback and questions! The responses to your questions are listed below, and paper is updated with updates colored in purple:

[W1] Relevance to previous work on EBM for autoregressive models.

Yes, [1] studied similar accumulated error problems (a.k.a exposure bias) in the autoregressive model setting and also proposed residual energy-based models for modeling global context. However, their methodology and analysis primarily focus on improving autoregressive models. By contrast, we study how to combine energy-based models with diffusion models, which is fundamentally different from theirs. Particular differences include:

1. Our Mask Diffusion EBM formulation inspires the idea of leveraging pretrained AR as EBM, while [1] can only on NCE training.
2. Even in NCE training, our framework involves new training objectives where positive and negative sampling distributions are from true and learned diffusion posteriors.
3. We introduce hybrid sampling combining both the denoising process and importance sampling.
4. We introduce the estimator of the variational lower bound of diffusion likelihood for evaluating the perplexity.

But following your suggestion, we realized we lack sufficient mention of this relevant work, and we have incorporated the content into related work.

[1] Deng, Yuntian, et al. "Residual energy-based models for text generation." arXiv preprint arXiv:2004.11714 (2020).

[W2] The description of the independent factorization in Equation 6.

In the paper, we use independent factorization in Equation 6 to represent that the distribution over the output distribution is factorized as a product of independent distributions where sampling is conducted independently.

Although the hidden representation is shared, this can be viewed as using shared parameters $\theta$ to model the independent distributions. This cannot model the correlations, since in the input space there is no information on other missing tokens. This argument is also used across other literatures [2].

[2] Lezama, Jose, et al. "Discrete predictor-corrector diffusion models for image synthesis." The Eleventh International Conference on Learning Representations. 2022.

[W3] Effectiveness on downstream tasks.

We follow your suggestion and test our EDLM on an additional semi-autoregressive generation setting. We follow exactly the same setting as MDLM Sec 4.2 [3], but change the sampling from diffusion denoising to energy-based importance sampling. In short, the model aims to generate sequences with 2048 tokens by autoregressively generating 512-length blocks. The Gen. PPL results are shown as follows:

As shown in the semi-autoregressive results, we also notice similar effectiveness translates to this downstream generation task.

[3] Sahoo, Subham Sekhar, et al. "Simple and Effective Masked Diffusion Language Models." arXiv preprint arXiv:2406.07524 (2024).

[W4] Additional inference cost.

Yes, our method will require more inference cost due to additional computation for importance sampling. However, we want to highlight that our method follows the current trend of language model acceleration research such as speculative decoding [4], where we use more computation to trade for shorter sampling time. Furthermore, with enough GPU memory, these computations can be conducted efficiently in parallel.

[4] Leviathan, Yaniv, Matan Kalman, and Yossi Matias. "Fast inference from transformers via speculative decoding, 2023." URL https://arxiv. org/abs/2211.17192 (2022).

Thank you for your constructive feedback and questions! The responses to your questions are listed below, and paper is updated with updates colored in purple:

[W1] This work primarily explores how to apply an energy-based framework to assist MLMs in generating text.

We agree with your takeaway that our method can be summarized as combining EBM and MLM through the diffusion framework for text generation. However, we highlight the key technical challenges and innovations as follows:

1. Our Mask Diffusion EBM framework inspires training with different noise levels determined by diffusion schedule, which differs from BERT.
2. We introduce NCE training for the EBM, which involves new training objectives where positive and negative sampling distributions are from true and learned diffusion posteriors.
3. We develop hybrid sampling combining both the denoising process and importance sampling.
4. We introduce the estimator of the variational lower bound of diffusion likelihood for evaluating the perplexity.

[W2] Using an AR model to rescore multiple outputs from an MLM is a standard baseline in non-autoregressive generation tasks

This is a good point. Though they also use the AR model, these methods typically use the AR to conduct ranking and selection to decode the most likely sample, where sampling methods are biased and cannot generate sampling according to learned distribution. By contrast, our EBM method can reduce sampling error and still induce unbiased sampling.

Theoretically, according to the formulation of mask diffusion models (Equation 5), the unmasking and masking schedules are defined in closed form, and the heuristic sampling methods such as ranking and selecting will result in biased sampling results. As a result, even if these methods can generate high-quality samples, they may result in other issues such as low diversity. As we can see, the relevant works primarily focus on conditional generation such as translation, where the metrics are mainly about quality instead of distribution modeling.

Empirically, furthermore, we also included SUNDAE [2] as the representation of these models in comparison, which is one of the SOTA models for translation. As shown in Tab 3 in the paper, SUNDAE achieves good results with an extremely small number of timesteps, but the diversity is obviously much worse than EBM methods. This indicates the biased sampling issue from these methods.

[1] Non-autoregressive Neural Machine Translation, 2017

[2] Savinov, Nikolay, et al. "Step-unrolled denoising autoencoders for text generation." arXiv preprint arXiv:2112.06749 (2021).

[Q1] Explanation of how to estimate PPL

Since perplexity is calculated from likelihood, we will mainly explain the details of calculating the likelihood. The exact formulation is in Eq. 13. Given data $x_0$, for every timestep $t$ we will perturb it to $x_t$ and compute the denoising loss (term 1 in Eq 13). For the energy term, we will first calculate the energy of $(x_0, x_t)$ pair (term 2); besides, we will sample several generated $x_0$ by denoising $x_t$ using the diffusion model and compute the average energy (this averaging is just an intuitive explanation for Theorem 1) of those samples to approximate the partition function (term 3). All three are combined together to compute Eq 13 weighted by the weights $\frac{\alpha_t’}{1-\alpha_t}$.

We hope our response could address your questions!

Dear reviewer,

This is a quick message to let you know that we are working on incorporating the suggested related work on Non-autoregressive Neural Machine Translation into the paper. We are still polishing the content and rearranging the pages, and will post here once we finished adding the content. We expect to finish it by tomorrow.

Thanks again for your feedback! We have updated the paper accordingly. Please let us know if you have other questions!

Thank you for your constructive feedback and questions! The responses to your questions are listed below, and paper is updated with updates colored in purple:

[W1] Relevance to previous work on EBM for autoregressive models.

Yes, [1] studied similar accumulated error problems (a.k.a exposure bias) in the autoregressive model setting and also proposed residual energy-based models for modeling global context. However, their methodology and analysis primarily focus on improving autoregressive models. By contrast, we study how to combine energy-based models with diffusion models, which is fundamentally different from theirs. Particular differences include:

1. Our Mask Diffusion EBM formulation inspires the idea of leveraging pretrained AR as EBM, while [1] can only rely on NCE training.
2. Even in NCE training, our framework involves new training objectives where positive and negative sampling distributions are from true and learned diffusion posteriors.
3. We introduce hybrid sampling combining both the denoising process and importance sampling.
4. We introduce the estimator of the variational lower bound of diffusion likelihood for evaluating the perplexity.

But following your suggestion, we realized we lack sufficient mention of this relevant work, and we have incorporated the content into related work.

[1] Deng, Yuntian, et al. "Residual energy-based models for text generation." arXiv preprint arXiv:2004.11714 (2020).

[W2] Discussion about relevant methods for filtering and remasking tokens at each denoising step.

This is a good point. We agree that diffusion provides likelihoods for all tokens and we can use this information to potentially identify the most incorrect tokens, but these sampling methods are biased and cannot generate sampling according to learned distribution. By contrast, our EBM method can reduce sampling error and still induce unbiased sampling.

Theoretically, according to the formulation of mask diffusion models (Equation 5), the unmasking and masking schedules are defined in closed form, and the heuristic sampling methods such as deterministic argmax decoding [2] will result in biased sampling results. As a result, even if these methods can generate high-quality samples, they result in other issues such as low diversity. As we can see, the relevant works [2-4] primarily focus on conditional generation such as translation or paraphrasing, where the metrics are mainly about quality instead of distribution modeling.

Empirically, furthermore, we also included SUNDAE [3] as the representation of these models in our benchmark comparison. As shown in Tab 3 in the paper, SUNDAE achieves good results with an extremely small number of timesteps, but the diversity is obviously much worse than EBM methods. This indicates the highly biased sampling issue from these methods.

[2] Ghazvininejad, Marjan, et al. "Mask-predict: Parallel decoding of conditional masked language models." arXiv preprint arXiv:1904.09324 (2019).

[3] Savinov, Nikolay, et al. "Step-unrolled denoising autoencoders for text generation." arXiv preprint arXiv:2112.06749 (2021).

[4] Zheng, Lin, et al. "A reparameterized discrete diffusion model for text generation." arXiv preprint arXiv:2302.05737 (2023).

[Q1] Details of pretrained AR model.

It’s a GPT-2 small model also trained on OpenWebText data, which uses the same model scale and training data as diffusion models. It’s a 12-layer transformer with 768-dimensional embeddings and 12 attention heads, resulting in 124M parameters.

[Q2] Baseline with an equivalent parameter count.

This is a good point. EDLM is composed of a diffusion model and an energy-based model, which introduces $2\times$ number of parameters. To clearly show that the performance improvements come from the framework design instead of model size, we follow your insight and conduct another ablation setting, where we double the number of layers and train a new MDLM for ~7 days. The key results on PPL are shown below:

As shown in the results, we didn’t notice a significant improvement by simply doubling the model size. This yields a more fair comparison to suggest the effectiveness of our proposed EDLM framework.

We hope our response could address your questions!

Thank you for your constructive feedback and questions! The responses to your questions are listed below, and paper is updated with updates colored in purple:

[W1] Unfair to compare with MDLM.

This is a good point. EDLM is composed of a diffusion model and an energy-based model, which introduces $2\times$ number of parameters. To clearly show that the performance improvements come from the framework design instead of model size, we follow your insight and conduct another ablation setting, where we double the number of layers and train a new MDLM for ~7 days. The key results on PPL are shown below:

As shown in the results, we didn’t notice a significant improvement by simply doubling the model size. While the neural network may become more expressive, the original model cannot account for correlations regardless of model size, which may explain the lack of a significant improvement. This yields a more fair comparison to suggest the effectiveness of our proposed EDLM framework.

[W2] GenPPL vs diversity trade-off is similar to MDLM.

We actually agree with your observation that the GenPPL vs diversity is similar. Note that, as in Fig1, MDLM’s diversity will also decrease with increased sample steps. EDLM aims to reduce the sampling errors and therefore has a similar performance as MDLM but with much fewer sampling timesteps.

[R1] More detailed trade-off between entropy and GenPPL.

Thank you for your suggestion and we prepared a figure for this comparison. We added it in Fig 4 in the updated paper. The figure supports our augment above, where we actually show a tradeoff similar to MDLM but with reduced time steps.

[R2] Improve efficiency for training the energy function $\phi$.

This is a good point and actually, in the EDLM-NCE setting, we train the energy $\phi$ by finetuning from the pretrained diffusion model $\theta$. We mention this information in Sec 5.1 Implementation Details paragraph: we finetune the pretrained MDLM, with the energy function computed by projecting the mean-pooled last token embeddings down to a single scalar value.

Besides, we would also like to highlight that in the other $\phi$ implementation EDLM-AR, we can leverage pretrained AR LLM and avoid additional training costs, which is even more efficient.

[Q1] Memory requirements for generation.

Yes, we approximately double the number of parameters. In FP-32 precision, the VRAM for MDLM is around 500MB while EDLM is 1GB. However, we emphasize that the VRAM during inference mainly comes from activation instead of the model itself, so this is not a bottleneck of EDLM.

[Q2] Importance sampling at different intervals.

Yes we can use different intervals. But as noted in Sec 4.4 Sampling window paragraph, we notice that with the same window, applying it on early stage sampling would always have better performance. For example, with 1024 sampling steps and a fixed window of 0.2, when applying it at different intervals we have the following results:

which shows that applying it at the start still shows the best result. We think the reason is the model tends to make more errors at the start due to limited context and importance sampling will correct them. Therefore we focus on discussing this case in the paper.

We hope our response could address your questions!