Edit-Based Refinement for Parallel Masked Diffusion Language Models

TL;DR

This paper introduces ME-DLM, an edit-based refinement framework for parallel masked diffusion language models that enhances multi-token generation quality and efficiency through minimal post-editing steps.

Contribution

It proposes a novel edit-based refinement method that improves sequence-level consistency and robustness in parallel diffusion language models.

Findings

Achieves 11.6 points improvement on HumanEval

Achieves 33.6 points improvement on GSM8K

Uses one-eighth of the diffusion steps for comparable performance

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.