kekchpek/idlm-duo-tynigsm

IDLM-Duo TinyGSM

IDLM-Duo TinyGSM is an Inverse-distilled Diffusion Language Model distilled from a pretrained Duo-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: Duo TinyGSM checkpoint from the S-FLM repository
• Diffusion type: uniform-state / Duo-style 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,153 tokens
• Parameters: 167,928,865
• 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-duo-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-duo-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.

Model	Sampling steps	Accuracy (%)
Duo teacher	1024	17.2
IDLM-Duo	32	15.39
IDLM-Duo	64	19.03
IDLM-Duo	128	21.38

Training Summary

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

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

The Duo setting uses a Gaussian relaxation and soft token inputs 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}
}