kekchpek/idlm-mdlm-tynigsm

IDLM-MDLM TinyGSM

IDLM-MDLM TinyGSM is an Inverse-distilled Diffusion Language Model distilled from a pretrained MDLM-style TinyGSM teacher checkpoint. It is released with the paper IDLM: Inverse-distilled Diffusion Language Models.

IDLM extends inverse distillation to discrete token spaces. Instead of running a pretrained diffusion language model for many reverse-diffusion steps, IDLM trains a few-step student generator using an auxiliary fake model and the teacher diffusion objective. This checkpoint targets conditional mathematical reasoning on TinyGSM-style question-answer examples.

• —Project page: https://david-cripto.github.io/idlm-project-page/
• —Code: https://github.com/David-cripto/IDLM
• —Paper: https://arxiv.org/abs/2602.19066
• —TinyGSM dataset: https://huggingface.co/datasets/TinyGSM/TinyGSM
• —Teacher/evaluation reference: https://github.com/jdeschena/s-flm

Model Details

• —Model family: IDLM, discrete diffusion language model
• —Teacher checkpoint: MDLM TinyGSM checkpoint from the S-FLM repository
• —Diffusion type: absorbing-state / MDLM-style masked diffusion
• —Training data: TinyGSM conditional question-answer examples
• —Tokenizer: HuggingFaceTB/SmolLM-135M tokenizer with an added [PAD] token
• —Context length: 512 tokens
• —Vocabulary size: 49,154 tokens
• —Parameters: 167,930,402
• —Tensor type: F32 Safetensors
• —Architecture config: 12 blocks, 12 heads, hidden size 768, conditioning dimension 128, dropout 0.1
• —Checkpoint files: model.safetensors contains the EMA backbone used for evaluation; idlm-mdlm-tynigsm.ckpt contains the PyTorch Lightning checkpoint

Intended Use

This checkpoint is intended for research on diffusion language models, inverse distillation, and few-step sampling.

Installation

The IDLM sampling and evaluation code depends on CUDA and FlashAttention.

git clone https://github.com/David-cripto/IDLM.git
cd IDLM

conda create -n idlm python=3.12
conda activate idlm
conda install nvidia/label/cuda-12.4.0::cuda-toolkit
pip install -r requirements.txt
pip install flash_attn==2.7.4.post1

Loading the Lightning Checkpoint

The full Lightning checkpoint can be inspected directly:

import torch

checkpoint = torch.load(
    "idlm-mdlm-tynigsm.ckpt",
    map_location="cpu",
)
print(checkpoint.keys())
print(next(iter(checkpoint["state_dict"])))

Evaluation

We report performance on conditional generation on GSM8K using the TinyGSM/GSM8K protocol from S-FLM. Each problem is sampled once, the generated solution is scored with the execution-based TinyGSM/GSM8K scorer, and the reported metric is exact-match accuracy against the GSM8K ground-truth answer.

Training Summary

IDLM-MDLM was trained by initializing the student and fake model from the pretrained MDLM teacher and alternating between:

1. 1.Updating the fake model on student-generated samples using the teacher diffusion loss.
2. 2.Updating the student using the teacher-fake loss gap.

The MDLM setting uses a mask-only simplex relaxation for stable backpropagation through the diffusion objective.

Citation

@article{li2026idlm,
  title={IDLM: Inverse-distilled Diffusion Language Models},
  author={Li, David and Gushchin, Nikita and Abulkhanov, Dmitry and Moulines, Eric and Oseledets, Ivan and Panov, Maxim and Korotin, Alexander},
  journal={arXiv preprint arXiv:2602.19066},
  year={2026}
}