RemeDi: Remasking-enabled Diffusion Language Model

🔬 Method Overview

RemeDi lets every token be revised at every diffusion step. Instead of fixing in an early guess, the model evaluates the quality of each token and can remask low-confidence positions, allowing later steps to resample them with richer context—built-in self-correction.

RemeDi extends the original model with a dual-stream transformer:

• Token Prediction Stream (TPS) predicts masked tokens as usual.

• Unmasking Policy Stream (UPS) outputs per-token confidence scores, deciding which tokens to unmask or remask.

At each denoising step, tokens with low confidence can be remasked and resampled, enabling iterative refinement. For the training and RL algorithms, see the Methods section of the paper.

📈 Key Results

🚀 Inference

To run inference, execute:

git clone https://github.com/maple-research-lab/RemeDi.git
cd RemeDi

# chat with remedi
python inference.py

📥 Citation

@article{huang2025don,
  title={Don't Settle Too Early: Self-Reflective Remasking for Diffusion Language Models},
  author={Huang, Zemin and Wang, Yuhang and Chen, Zhiyang and Qi, Guo-Jun},
  journal={arXiv preprint arXiv:2509.23653},
  year={2025}
}