Encoder-Decoder Diffusion Language Models for Efficient Training and Inference

We thank the reviewer for their feedback. Please see below for responses to the specific criticisms and questions raised.

---

Concern 1: Limited novelty

We argue that our work has novelty beyond combining a T5-style architecture with BD3-LM. Specifically, it involves the following research challenges:

• Architecture Design. There are multiple ways to define encoder-decoder inputs and outputs. See e.g., ENAT [5] for an alternative architecture optimized for images. More generally, we report ablations on ways of connecting the encoder and decoder, the interplay of block size and decoder size, choices of decoder optimizations, weight-tying, and more.
• Training Algorithms. Our architecture requires specialized training algorithms that extend BD3-LM (Algorithm 1), which rely on custom attention masks relative to BD3LM (Figure 2) and a modified computation graph (Figure 1) relative to T5.
• Sampling Algorithms. Similarly, we introduce sampling algorithms that extend BD3-LM (Algorithm 2) and provide hyper-parameters (e.g., block size, number of layers) needed to obtain good performance.

---

Concern 2: Clarifying why fine-tuned models have drop in performance relative to base model

There are several factors at play here. First, although E2D2 is within 1 perplexity point of the base Qwen model [1], this gap translates to a larger drop in GSM8K exact-match accuracy which is sensitive to small mistakes in the arithmetic trace. Qualitatively, we observe that even single-token mismatches may result in an incorrect response. Additionally, the gap between even the fine-tuned AR model and the value reported in the original Qwen3 technical report [1] stems from the different evaluation setup, as detailed below.

GSM8k Results: (Best diffusion values are bolded)

(Note: the E2D2 values here represent an improved fine-tuning setup with smaller batch size, which we will update in our revised manuscript).

In terms of the differences in evaluation setup:

• The original Qwen model used 4-shot CoT reasoning traces. Our evaluation is 0-shot.
• Additionally, we fine-tuned the model to just output “reasoning traces” provided in the GSM8k dataset. This differs significantly from the default output that the Qwen model produces which is much more verbose. We limited our scope of number of shots and reasoning traces for computational considerations.

We also comment that this gap is limited to the GSM8k dataset. On the translation and summarization tasks, however, E2D2 beats the performance of the fine-tuned AR model (Tables 5 & 6):

---

Concern 3: Limited improvement of E2D2 vs. BD3LM

While performance relative to BD3LM is comparable in Table 5, we note that we achieve this at either improved efficiency (relative to the train-matched BD3LM) or no additional cost (relative to the inference matched BD3LM).

To make this tradeoff of efficiency and performance more concrete, we provide a table showing the Pareto frontier between model performance (measured by PPL and pass@1) and throughput for E2D2 vs. BD3LM models. We find that for comparable throughput, E2D2 consistently features higher quality. Thus, it extends forward the Pareto frontier of quality and speed. We obtain best results for high and medium throughput levels, where the gain in quality is especially high.

To measure this frontier we vary model size (number of layers) which controls both the performance (larger model $\rightarrow$ higher performance) and throughput (larger model $\rightarrow$ lower throughput) axes.

Exp. details:

• We fine-tuned Qwen 1.7B-Base models on the gsm8k train dataset.
• Models were fine-tuned with batch size of 1, for 30k steps, which we found to be a more useful setup than the one currently reported in the manuscript. We will update Table 7 with these[ improved numbers .
• Throughput calculation was done using batch size of 1 on a single H100 GPU. Models generated 256 new tokens for 50 prompts from the gsm8k test dataset. We report mean $\pm$ standard deviation across samples. Note also that the throughput values reported here are superior to those in the original manuscript as we found inefficiencies in our code that we were able to streamline.
• For MDLM, we generate semi-autoregressively using blocks of 64, to have a fair comparison to the other models. We did a sweep on block sizes from 4 to 256 for MDLM to find the best performing one.
• Note, E2D2 has the same trainable parameters as AR and MDLM since we tie the weights of the models here.

---

Concern 4: Missing related works

Thank you for these suggested related works. We are happy to include them in our revised manuscript. Unfortunately, the exact references for “[3]” and “[4]” suggested by the reviewer appear to have been cut off in the OpenReview UI. We therefore present discussions for what we assume are the works to which the reviewer was referring. Please let us know if we are not discussing the ones you had in mind:

“multi-token prediction [3]”: Gloeckle, Fabian, et al. "Better & faster large language models via multi-token prediction." (2024).

Our work differs from [3] in several important ways. First, for their non-parallel generation, within a block, [3] maintains an autoregressive decoding scheme. In contrast, our diffusion parameterization within each block enables any-order generation. Additionally, similar to our contrast with works, such as Medusa [5], our architecture represents a more expressive decoder relative to the “heads” used in [3] and [5] and E2D2 has more sophisticated information sharing between the encoder and the decoder via cross-attention to the encoder outputs.

“next block generation [4]” Ho, Namgyu, et al. "Block transformer: Global-to-local language modeling for fast inference." (2024)

The key difference between E2D2 and this work are as follows: first [4] is still parameterized as an AR model, whereas E2D2 uses a block diffusion parameterization. Additionally, E2D2 employs standard global, self-attention whereas the efficiencies from [4] come from a custom hierarchical transformer model that first pools tokens to create block embeddings over which global attention is applied and then performs only local attention within blocks to decode autoregressively. These efficiencies are orthogonal to our proposed method and in fact could potentially be combined with E2D2 by replacing the standard transformer with the Block Transformer, to reduce KV cache memory storage and retrieval bottlenecks.

---

Question 1: Can E2D2’s encoder support efficient KV caching

We clarify that E2D2 also supports efficient caching. Both the encoder and decoder use a KV cache of the already decoded / clean context blocks. More specifically, in E2D2, the encoder first processes clean context tokens. These are stored in a KV cache for the encoder (which the decoder also uses). Each newly decoded block is added to this KV cache in the same manner they would for BD3LM, except in E2D2, this is done by processing them in the encoder.

Thus KV-cache is not a bottleneck for E2D2 models in scaling to multi-turn, long-dialogue settings.

---

Question 2: Clarifying the model sizes used in each experiment

Thank you for the suggestion. This information (reproduced in the table below) is included in the appendix, but to improve readability we will add it to the main text and corresponding tables as well.

---

References:

[3] Gloeckle, Fabian, et al. "Better & faster large language models via multi-token prediction." arXiv preprint arXiv:2404.19737 (2024).

[4] Ho, Namgyu, et al. "Block transformer: Global-to-local language modeling for fast inference." Advances in Neural Information Processing Systems 37 (2024): 48740-48783.

[5] Cai, Tianle, et al. "Medusa: Simple llm inference acceleration framework with multiple decoding heads." arXiv preprint arXiv:2401.10774 (2024).

Dear Reviewer,

Thank you again for providing your initial feedback on our work. Given that the discussion period ends in two days, we wanted to reach out and see if there were any additional points of clarification we can provide. Did our rebuttal address your concerns and questions? If so, would you kindly be willing to reconsider your initial score of our work?

Looking forward to your feedback.

Sincerely,

The Authors

We thank the reviewer for their valuable feedback and for recognizing the contributions of our work. Below we respond to the specific concerns and questions

---

Concern 1: Provide a more complete picture of the “pareto frontier”

Below, we provide a table showing the Pareto frontier between model performance (measured by PPL and pass@1) and throughput for E2D2 vs. BD3LM models. We find that for comparable throughput, E2D2 consistently features higher quality. Thus, it extends forward the Pareto frontier of quality and speed. We obtain best results for high and medium throughput levels, where the gain in quality is especially high.

To measure this frontier we vary model size (number of layers) which controls both the performance (larger model $\rightarrow$ higher performance) and throughput (larger model $\rightarrow$ lower throughput) axes.

Exp. details:

• Best values beyond margin of error are bolded
• We fine-tuned Qwen 1.7B-Base models on the gsm8k train dataset.
• Models were fine-tuned with batch size of 1, for 30k steps, which we found to be a more useful setup than the one currently reported in the manuscript. We will update Table 7 with these[ improved numbers .
• Throughput calculation was done using batch size of 1 on a single H100 GPU. Models generated 256 new tokens for 50 prompts from the gsm8k test dataset. We report mean $\pm$ standard deviation across samples. Note also that the throughput values reported here are superior to those in the original manuscript as we found inefficiencies in our code that we were able to streamline.
• For MDLM, we generate semi-autoregressively using blocks of 64, to have a fair comparison to the other models. We did a sweep on block sizes from 4 to 256 for MDLM to find the best performing one.
• Note, E2D2 has the same trainable parameters as AR and MDLM since we tie the weights of the models here.

---

Concern 2: Explore other design choices for decoupling Encoder-Decoder architecture design choices are underexplored

Here, we explore an alternative way of connecting the encoder and decoder: rather than sharing KV cache for matched layers in the encoder decoder, we feed the last hidden layer representation of the encoder into the decoder. We conducted an experiment for this new model by training from scratch on the machine translation (WMT14 de-en) dataset. In this setup, we also find that E2D2 is more efficient and better performing compared to MDLM and BD3LM.

WMT results:

Although E2D2 has more parameters than the BD3LM baseline, it is both faster and has better downstream performance, highlighting the benefit of our proposed encoder-decoder structure.

WMT Exp. details:

• We train models with hidden dimension 512 and MLP dimension of 1536 from scratch. Other model architecture details follow Qwen3 models.
• Batch size of 128, LR 3e-4, trained for 1M steps. We perform early stopping using eval loss.
• Throughput calculation was done using batch size of 1 on a single H100 GPU. Models generated 256 new tokens for 50 prompts from the test dataset. We report mean $\pm$ standard deviation across samples.
• For MDLM, we generate semi-autoregressively using blocks of 32, to have a fair comparison to the other models. We did a sweep on block sizes from 4 to 256 for MDLM to find the best performing one.
• Note: We use a modified architecture for E2D2 here: encoder and decoder weights are untied. Rather than sharing KV cache between the encoder and decoder as above, we pass the last hidden representation from the encoder to each layer of the decoder.

Next, we address this comment from the reviewer:

Another direction for exploration is whether the size of the decoder affects the number of steps that the diffusion model needs to perform for a fixed target perplexity.

Below we present an experiment that compares varying the number of diffusion steps per block for E2D2, across different decoder sizes. The efficiency of decoupling encoder and decoder compute in E2D2 allows us to use larger models with fewer decoding steps and still outperform BD3LM models (that use more diffusion steps) in terms of both the throughput and benchmark performance axes.

(Exp. details: Same as above)

---

Concern 3: Add a “from scratch” experiment

This comment is well received. The WMT experiment from above was indeed conducted from scratch and demonstrates that our approach is not reliant on initializing from pre trained models.

---

Formatting concerns:

Thank you for these comments. We will update the revised manuscript accordingly.

---

Question 1: “Is the use of a KV-cache novel to this work?”

To clarify, both E2D2 and BD3LM enable KV-caching of already decoded blocks and this is one of the primary motivations of building E2D2 on BD3LM models as opposed to their predecessor vanilla MDLM models, which do not support KV caching (in addition to the fact that BD3LM boasts improved performance vs. MDLM). The difference between these models is that to decode a given token one needs to make a pass through all the layers of the BD3LM model, whereas for E2D2 we only require a forward pass through the smaller number of decoder layers.

Thus, for the sake of illustration, assume that the E2D2 encoder is the same size as the full BD3LM model, then the KV cache for both BD3LM and E2D2 would be the same: cache any context/prompt tokens and, as each new block is generated, add it to the cache (by running it through the full network for BD3LM or through the encoder for E2D2). For E2D2, the decoder uses the KV cache from the encoder as well.

---

Question 2: “Could the smaller diffusion models make more sense the larger the encoder network is?”

Great question. We explore this above in the response to concern 1.

---

Question 3: “How were the ‘inference-matched’ and ‘train-matched’ baselines chosen?”

These baselines were chosen based on the FLOPs computations provided in Tables 1 and 2.

• Train-matched: To match the training flops between E2D2 and BD3LM, the BD3LM model needs to have half as many layers as the combined number of encoder and decoder layers of the E2D2 model. See Table 1.
• Inference-matched: We matched the inference compute empirically by seeing the number of layers of a BD3LM model that would lead to similar throughput to an E2D2 model (assuming we only decode one token per forward pass of the diffusion model, i.e. diffusion time steps $T$ per block is equal to block size $S$). See the throughput results from the response to Concern 1 above.

Dear Reviewer,

Thank you again for providing your initial feedback on our work. Given that the discussion period ends in two days, we wanted to reach out and see if there were any additional points of clarification we can provide. Did our rebuttal address your concerns and questions? If so, would you kindly be willing to reconsider your initial score of our work?

Looking forward to your feedback.

Sincerely,

The Authors

Thank you again for the thoughtful feedback and discussion, which were instrumental in helping us gather and better frame these new results. The new results will certainly be included in our updated manuscript.

We hope that this positive and productive discussion period will be reflected in your final assessment of our work.

Sincerely,

The Authors

We thank the reviewer for their feedback and for the valuable additional experiment suggestions. Below we respond to the specific comments inline.

---

Concern 1: Provide a more complete “pareto frontier” picture

Per the reviewer request, below, we provide a table showing the Pareto frontier between model performance (measured by PPL and pass@1) and throughput for E2D2 vs. BD3LM models. We find that for comparable throughput, E2D2 consistently features higher quality. Thus, it extends forward the Pareto frontier of quality and speed. We obtain best results for high and medium throughput levels, where the gain in quality is especially high.

To measure this frontier we vary model size (number of layers) which controls both the performance (larger model $\rightarrow$ higher performance) and throughput (larger model $\rightarrow$ lower throughput) axes.

(Note: This table also address the reviewer’s request for presenting the “unmodified” BD3LM, i.e., the 28 layer BD3LM).

Exp. details:

• Best values beyond margin of error are bolded
• We fine-tuned Qwen 1.7B-Base models on the gsm8k train dataset.
• Models were fine-tuned with batch size of 1, for 30k steps, which we found to be a more useful setup than the one currently reported in the manuscript. We will update Table 7 with these[ improved numbers .
• Throughput calculation was done using batch size of 1 on a single H100 GPU. Models generated 256 new tokens for 50 prompts from the gsm8k test dataset. We report mean $\pm$ standard deviation across samples. Note also that the throughput values reported here are superior to those in the original manuscript as we found inefficiencies in our code that we were able to streamline.
• For MDLM, we generate semi-autoregressively using blocks of 64, to have a fair comparison to the other models. We did a sweep on block sizes from 4 to 256 for MDLM to find the best performing one.
• Note, E2D2 has the same trainable parameters as AR and MDLM since we tie the weights of the models here.

---

Concern 2: Explore the effect of block size more

Below we present the exploration of the effect of block size for E2D2 models. We find that E2D2 with block size of either 4 or 8 better trades-off benchmark performance with throughput.

Note that all models above were fine-tuned as separate runs (see the answer to the related question below). This is an updated version of Table 10 from our paper which we will revise in our latest manuscript.

Exp details: Same as above.

---

Question 1: How is the AR model trained / what does the model size look like?

All of the models are based on the same backbone and are hence comparable. The AR model in particular represents an unmodified backbone, that is we use all 28 layers of the pre-trained model. This is akin to matching the number of encoder layers of E2D2. This renders the AR model a very strong baseline. We will add this information more explicitly in our revised manuscript

---

Question 2: Are the experiments in 4.4 fine tuned individually (i.e., for each parameter) or the results with the same base model?

We apologize for the confusion. In our updated manuscript we will make it clear that each result in the ablations section (4.4) for both Table 9 and Table 10 represents a separate fine-tuning run. That is, for each row in these tables, starting from the same baseline pre-trained model, we initialize weights from the pre-trained model, modify the number or layers (Table 9) or block size (Table 10), and then fine-tune. This holds true for the updated results presented in this response as well, which we will use when revising our manuscript.

Dear Reviewer,

Thank you again for providing your initial feedback on our work. Given that the discussion period ends in two days, we wanted to reach out and see if there were any additional points of clarification we can provide. Did our rebuttal address your concerns and questions? If so, would you kindly be willing to reconsider your initial score of our work?

Looking forward to your feedback.

Sincerely,

The Authors

We thank the reviewer for their feedback and for recognizing the contribution of E2D2 in improving discrete diffusion models. Below we respond to the reviewer’s concerns and questions.

Concern 1: Missing comparison to models such as MDLM

We thank the reviewer for this suggestion. Below we provide this comparison on the GSM8k dataset. We find that MDLM significantly underperforms E2D2 on this task.

GSM8k Results: (Best diffusion values are bolded)

We also provide this comparison for models trained from scratch on the machine translation (WMT de-en) dataset. Here too, E2D2 is more efficient and better performing.

WMT results:

Although E2D2 has more parameters than the BD3LM baseline, it is both faster and has better downstream performance, highlighting the benefit of our proposed encoder-decoder structure.

GSM8k Exp. details:

• We fine-tuned Qwen 1.7B-Base models on the gsm8k train dataset.
• Models were fine-tuned with batch size of 1, for 30k steps, which we found to be a more useful setup than the one currently reported in the manuscript. We will update Table 7 with these improved numbers .
• Throughput calculation was done using batch size of 1 on a single H100 GPU. Models generated 256 new tokens for 50 prompts from the gsm8k test dataset. We report mean $\pm$ standard deviation across samples. Note also that the throughput values reported here are superior to those in the original manuscript as we found inefficiencies in our code that we were able to streamline.
• For MDLM, we generate semi-autoregressively using blocks of 64, to have a fair comparison to the other models. We did a sweep on block sizes from 4 to 256 for MDLM to find the best performing one.
• Note, E2D2 has the same trainable parameters as AR and MDLM since we tie the weights of the models here.

WMT Exp. details:

• We train models with hidden dimension 512 and MLP dimension of 1536 from scratch. Other model architecture details follow Qwen3 models.
• Batch size of 128, LR 3e-4, trained for 1M steps. We perform early stopping using eval loss.
• Throughput calculation was done using batch size of 1 on a single H100 GPU. Models generated 256 new tokens for 50 prompts from the test dataset. We report mean $\pm$ standard deviation across samples.
• For MDLM, we generate semi-autoregressively using blocks of 32, to have a fair comparison to the other models. We did a sweep on block sizes from 4 to 256 for MDLM to find the best performing one.
• Note: We use a modified architecture for E2D2 here: encoder and decoder weights are untied. Rather than sharing KV cache between the encoder and decoder as above, we pass the last hidden representation from the encoder to each layer of the decoder.

---

Concern 2: Explaining the relationship between decoding throughput and number of decoder layers and block size

As demonstrated in the asymptotic FLOPs computation of Tables 1 and 2, training and inference computational budgets are functions of both the number of layers and the block size. This is true for both E2D2 and BD3LM (and true for all models, e.g., AR and MDLM, in terms of the number of layers).

In terms of the reviewer’s question on this topic:

Regarding weakness 2, how can the optimal settings for the number of decoder layers and block size be determined?

The “optimal” architecture balances downstream performance and efficiency (i.e. throughput), which can be user-dependent. Models that better control this trade-off are preferable.

Below, we provide a table showing the Pareto frontier between model performance (measured by PPL and pass@1) and throughput for E2D2 vs. BD3LM models. We find that for comparable throughput, E2D2 consistently features higher quality. Thus, it extends forward the Pareto frontier of quality and speed. We obtain best results for high and medium throughput levels, where the gain in quality is especially high.

To measure this frontier we vary model size (number of layers) which controls both the performance (larger model $\rightarrow$ higher performance) and throughput (larger model $\rightarrow$ lower throughput) axes.

(Exp. details: Same as above; Best values beyond margin of error are bolded)

---

Concern 3: Missing discussion of Llada and predecessor models (MDLM)

We thank the reviewer for this suggestion, and we will correct this oversight in the updated manuscript. Specifically we edit line 275 in the “Diffusion Language Models” paragraph of the related works as follows:

Recent large-scale diffusion language models have demonstrated competitive performance with comparably sized AR LLMs on benchmark metrics. Specifically, the Llada model [2], built on the work of MDLM [3] Nie et al., 2024 [4], scales masked diffusion models to the 8B parameter regime and demonstrated comparable and superior performance to a similarly-sized AR Llama model [5].

---

Concern 4: Missing “limitations” section

Thank you for this suggestion. We will add the following text as a limitations section:

While our work represents a further step in closing the performance and efficiency gap to the dominant AR paradigm for language modeling, our experimental results demonstrate that discrete diffusion models still lag. Further innovation in training and inference algorithms for discrete diffusion models is needed. Additionally, since our work entails language modeling it carries the inherent risks of misuse of this technology. As these models improve in size and quality, care must be taken to prevent malicious use.

---

Additional concerns: Table placement and typos

Thank you for these suggestions. We will update the formatting and text in our revised manuscript.

---

Question: Clarify the primary advantages of diffusion relative to AR, which has faster throughput

As noted in the updated table above, we have removed an important inefficiency in our code which brings the throughput of the diffusion models in line with that of the AR baseline. Moveover, we see the promise of E2D2 which is able to have the best diffusion performance with faster throughput than AR models. We apologize that these results were not reflected in our original manuscript, but we will update them in our revised version.

---

References

[1] Arriola, Marianne, et al. “Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models.” ICLR 2025

[2] Nie, Shen, et al. "Large language diffusion models." arXiv preprint arXiv:2502.09992 (2025).

[3] Sahoo, Subham, et al. "Simple and effective masked diffusion language models." Advances in Neural Information Processing Systems 37 (2024): 130136-130184.

[4] Nie, Shen, et al. "Scaling up masked diffusion models on text." arXiv preprint arXiv:2410.18514 (2024).

[5] Dubey, Abhimanyu, et al. "The llama 3 herd of models." arXiv e-prints (2024): arXiv-2407.