Edit-Based Refinement for Parallel Masked Diffusion Language Models
Abstract
Masked diffusion language models enable parallel token generation and offer improved decoding efficiency over autoregressive models. However, their performance degrades significantly when generating multiple tokens simultaneously, due to a mismatch between token-level training objectives and joint sequence consistency. In this paper, we propose ME-DLM, an edit-based refinement framework that augments diffusion generation with lightweight post-editing steps. After producing an initial complete response, the model refines it through minimal edit operations, including replacement, deletion, and insertion, conditioned on the full sequence. Training supervision is derived from edit distance, providing a deterministic signal under a fixed canonicalization scheme for learning minimal corrections. This approach encourages sequence-level consistency through globally conditioned edits while preserving the efficiency benefits of parallel diffusion decoding. Extensive experiments demonstrate that ME-DLM improves the quality and robustness of multi-token parallel generation. In particular, when built upon LLaDA, our method achieves consistent gains of 11.6 points on HumanEval and 33.6 points on GSM8K while using one-eighth of the total diffusion steps. Code is available at https://github.com/renhouxing/ME-DLM.
Lay Summary
Large language models usually generate text one word at a time, which can be slow. An alternative approach, called diffusion language modeling, can generate many words in parallel by gradually filling in missing parts of a sentence. This is faster, but it often becomes less reliable when many words are generated at once, because the model may choose words that look reasonable individually but do not fit well together as a complete answer. We introduce ME-DLM, a method that helps diffusion language models improve their own answers after an initial draft is produced. Instead of generating everything again, the model makes small edits—such as replacing, deleting, or inserting words—while looking at the whole response. This allows the model to fix inconsistencies that are hard to catch when words are predicted separately. Our experiments show that this simple editing step makes parallel generation both more accurate and more robust. Built on top of LLaDA, ME-DLM substantially improves performance on coding and math benchmarks while using far fewer diffusion steps, helping preserve the speed advantage of diffusion-based text generation.