Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Teach Diffusion Language Models to Learn from Their Own Mistakes

Liming Liu, Binxuan Huang, Xin Liu, Bing Yin, Tuo Zhao

TL;DR

This work tackles the vulnerability of Masked Diffusion Language Models (MDMs) to dependency errors during fast, parallel decoding. It proposes Decoupled Self-Correction (DSC), a two-stage training scheme that preserves base generative fidelity by first fully training the generator and then training a lightweight correction head, and Future-Context Augmentation (FCA) to expose the head to richer future contexts. The approach enables reliable self-correction at inference via adaptive remasking, significantly mitigating error accumulation and improving generation quality on coding and mathematical reasoning benchmarks, while maintaining speed advantages. The results show that DSC expands the Pareto frontier of speed versus fidelity and demonstrates robustness across domains, with ablations validating the importance of aligned artifact distributions and future context.

Abstract

Masked Diffusion Language Models (DLMs) achieve significant speed by generating multiple tokens in parallel. However, this parallel sampling approach, especially when using fewer inference steps, will introduce strong dependency errors and cause quality to deteriorate rapidly as the generation step size grows. As a result, reliable self-correction becomes essential for maintaining high-quality multi-token generation. To address this, we propose Decoupled Self-Correction (DSC), a novel two-stage methodology. DSC first fully optimizes the DLM's generative ability before freezing the model and training a specialized correction head. This decoupling preserves the model's peak SFT performance and ensures the generated errors used for correction head training are of higher quality. Additionally, we introduce Future-Context Augmentation (FCA) to maximize the correction head's accuracy. FCA generalizes the error training distribution by augmenting samples with ground-truth tokens, effectively training the head to utilize a richer, future-looking context. This mechanism is used for reliably detecting the subtle errors of the high-fidelity base model. Our DSC framework enables the model, at inference time, to jointly generate and revise tokens, thereby correcting errors introduced by multi-token generation and mitigating error accumulation across steps. Experiments on mathematical reasoning and code generation benchmarks demonstrate that our approach substantially reduces the quality degradation associated with larger generation steps, allowing DLMs to achieve both high generation speed and strong output fidelity.

Teach Diffusion Language Models to Learn from Their Own Mistakes

TL;DR

This work tackles the vulnerability of Masked Diffusion Language Models (MDMs) to dependency errors during fast, parallel decoding. It proposes Decoupled Self-Correction (DSC), a two-stage training scheme that preserves base generative fidelity by first fully training the generator and then training a lightweight correction head, and Future-Context Augmentation (FCA) to expose the head to richer future contexts. The approach enables reliable self-correction at inference via adaptive remasking, significantly mitigating error accumulation and improving generation quality on coding and mathematical reasoning benchmarks, while maintaining speed advantages. The results show that DSC expands the Pareto frontier of speed versus fidelity and demonstrates robustness across domains, with ablations validating the importance of aligned artifact distributions and future context.

Abstract

Masked Diffusion Language Models (DLMs) achieve significant speed by generating multiple tokens in parallel. However, this parallel sampling approach, especially when using fewer inference steps, will introduce strong dependency errors and cause quality to deteriorate rapidly as the generation step size grows. As a result, reliable self-correction becomes essential for maintaining high-quality multi-token generation. To address this, we propose Decoupled Self-Correction (DSC), a novel two-stage methodology. DSC first fully optimizes the DLM's generative ability before freezing the model and training a specialized correction head. This decoupling preserves the model's peak SFT performance and ensures the generated errors used for correction head training are of higher quality. Additionally, we introduce Future-Context Augmentation (FCA) to maximize the correction head's accuracy. FCA generalizes the error training distribution by augmenting samples with ground-truth tokens, effectively training the head to utilize a richer, future-looking context. This mechanism is used for reliably detecting the subtle errors of the high-fidelity base model. Our DSC framework enables the model, at inference time, to jointly generate and revise tokens, thereby correcting errors introduced by multi-token generation and mitigating error accumulation across steps. Experiments on mathematical reasoning and code generation benchmarks demonstrate that our approach substantially reduces the quality degradation associated with larger generation steps, allowing DLMs to achieve both high generation speed and strong output fidelity.
Paper Structure (32 sections, 8 equations, 3 figures, 7 tables, 1 algorithm)

This paper contains 32 sections, 8 equations, 3 figures, 7 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Our Contribution
  3. Concurrent Work
  4. Related Works
  5. Preliminaries
  6. Masked Diffusion Models
  7. Forward process.
  8. Reverse process.
  9. Training.
  10. Inference.
  11. Self-correction of MDMs
  12. Method
  13. Decoupled Self-Correction (DSC) Training
  14. Stage 1: Optimized Generative Training:
  15. Stage 2: Specialized Head Training:
  16. ...and 17 more sections

Figures (3)

  • Figure 1: The accuracy vs. iterations comparison between our method and baseline SFT on GSM8k and MBPP. Data is from Table \ref{['tab:mbpp_results']} and Table \ref{['tab:gsm8k_results']}. Our methods is Pareto better than baseline.
  • Figure 2: The accuracy vs. iterations comparison between our method and the ablation method in \ref{['sec:ablation']} on GSM8k and MBPP. Data is from \ref{['tab:gsm8k_results_ablation']}. Our methods is Pareto better than the ablation method.
  • Figure 3: The accuracy vs. iterations comparison between our method and randomly remasking method in \ref{['sec:random']} on GSM8k. Data is from \ref{['tab:gsm8k_results_random']}. Our methods is Pareto better than the ablation method.