IDLM-MDLM
IDLM-MDLM is an Inverse-distilled Diffusion Language Model distilled from the pretrained
MDLM OpenWebText checkpoint. It is released with the paper
IDLM: Inverse-distilled Diffusion Language Models.
Diffusion Language Models can produce high-quality text, but standard reverse diffusion requires many sampling steps. IDLM trains a few-step student generator from a pretrained DLM teacher using an inverse distillation objective with an auxiliary fake model. This checkpoint targets fast generation from an absorbing-state masked diffusion teacher.
Model Details
- Model family: IDLM, discrete diffusion language model
- Teacher checkpoint:
kuleshov-group/mdlm-owt
- Diffusion type: absorbing-state / masked diffusion
- Training data: OpenWebText
- Tokenizer: GPT-2 tokenizer
- Context length: 1024 tokens
- Parameters: 169,627,250
- Tensor type: F32 Safetensors
- Architecture config: 12 blocks, 12 heads, hidden size 768, conditioning dimension 128, dropout 0.1
- License: MIT
Intended Use
This checkpoint is intended for research on discrete diffusion language models, few-step diffusion sampling, and reproduction of the IDLM paper experiments.
Installation
The sampling code depends on CUDA and FlashAttention.
1git clone https://github.com/David-cripto/IDLM.git
2cd IDLM
3
4conda create -n idlm python=3.12
5conda activate idlm
6conda install nvidia/label/cuda-12.4.0::cuda-toolkit
7pip install -r requirements.txt
8pip install flash_attn==2.7.4.post1
Loading the Checkpoint
The Hugging Face repository contains custom model code. Use trust_remote_code=True.
1from transformers import AutoModelForMaskedLM, AutoTokenizer
2
3model_id = "kekchpek/idlm-mdlm"
4tokenizer = AutoTokenizer.from_pretrained(model_id)
5model = AutoModelForMaskedLM.from_pretrained(
6 model_id,
7 trust_remote_code=True,
8)
Direct AutoModelForMaskedLM loading exposes the denoising network. For text generation, use the sampler in the official IDLM repository.
Generate Samples
1mkdir -p samples
2
3python -m main \
4 mode=sample_eval \
5 loader.batch_size=2 \
6 loader.eval_batch_size=8 \
7 data=openwebtext-split \
8 algo=mdlm \
9 algo.backbone=hf_dit \
10 eval.checkpoint_path=kekchpek/idlm-mdlm \
11 sampling.steps=16 \
12 sampling.num_sample_batches=10 \
13 sampling.predictor=ancestral_cache \
14 sampling.noise_removal=ancestral \
15 +wandb.offline=true \
16 eval.generated_samples_path=samples/idlm_mdlm_16steps.json
The generation script can be swept with different sampling steps.
Evaluation
The paper reports generation perplexity (GenPPL, lower is better) and sample entropy (higher is better) on OpenWebText-style generation. The released evaluation code defaults to gpt2-large for GenPPL.
| Sampling steps | GenPPL (lower is better) | Entropy (higher is better) |
|---|
| 32 | 20.37 | 5.23 |
| 16 | 32.74 | 5.42 |
| 8 | 79.42 | 5.61 |
| 4 | 310.38 | 5.78 |
For comparison, the MDLM teacher is reported at 1024 steps with GenPPL 41.29 and entropy 5.28.
Training Summary
IDLM-MDLM was trained by initializing the student and fake model from the pretrained MDLM teacher and alternating between:
- Updating the fake model on student-generated samples using the teacher diffusion loss.
- Updating the student using the teacher-fake loss gap.
This follows the inverse distillation objective described in the paper and uses the absorbing-state masked diffusion formulation.
Citation
1@article{li2026idlm,
2 title={IDLM: Inverse-distilled Diffusion Language Models},
3 author={Li, David and Gushchin, Nikita and Abulkhanov, Dmitry and Moulines, Eric and Oseledets, Ivan and Panov, Maxim and Korotin, Alexander},
4 journal={arXiv preprint arXiv:2602.19066},
5 year={2026}
6}