jh27kim/cdm-checkpoints

CDM — released checkpoints

Checkpoints for Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion (CDM).

• 📄 Paper: arXiv:2605.23346
• 🌐 Project page: https://cdm-smc.github.io/
• 💻 Code: https://github.com/KAIST-Visual-AI-Group/CDM

CDM trains a twist head psi_theta once with a contrastive objective so that twisted Sequential Monte Carlo no longer needs a Monte Carlo rollout at every denoising step. This repository holds the trained twist heads, plus the base models and reward oracles that are not available elsewhere on the Hub.

Contents

toxicity/
├── mdlm.ckpt                  # base MDLM (DiT, OpenWebText)
└── cdm/twist_best.pt          # trained twist head
dna/
├── mpra.ckpt                  # base MDLM (CNN, Gosai enhancers)
├── reward_oracle_ft.ckpt      # given reward: HepG2 Enformer oracle
├── reward_oracle_eval.ckpt    # heldout reward: second Enformer, validation split
├── human_state_dict.h5        # Enformer backbone weights (grelu artifact cache)
└── cdm/twist_best.pt          # trained twist head
proteins/
└── cdm/twist_best.pt          # trained twist head
dllm/
└── cdm/twist_best.pt          # trained twist head

Everything else the code needs is pulled from the Hub on first use: DPLM-2 (airkingbd/dplm2_650m), ESMFold, the two RoBERTa toxicity classifiers, Skywork-Reward-Llama-3.1-8B and ArmoRM-Llama3-8B.

The diffusion-LLM base model is LLaDA-8B-Instruct; an HF-format copy is available at MSALab/LLaDA-8B-Instruct-HF.

Usage

Clone the code first — the checkpoints are only useful alongside it:

git clone https://github.com/KAIST-Visual-AI-Group/CDM
cd CDM

Then download everything straight into place:

python scripts/download_checkpoints.py                # all four applications
python scripts/download_checkpoints.py --apps dna     # just one

That script puts each file where the configs expect it, under cdm/<app>/checkpoints/. To fetch manually instead:

from huggingface_hub import hf_hub_download

hf_hub_download("jh27kim/cdm-checkpoints", "dna/cdm/twist_best.pt",
                local_dir="cdm/dna/checkpoints_raw")

Sampling with a trained twist, once the files are in place:

python -m cdm.texts_mdm.main --config-name cdm K=8 \
  twist_ckpt=./cdm/texts_mdm/checkpoints/cdm/twist_best.pt

See the repository README for environment setup and the full command reference.

Citation

@article{kim2026cdm,
  title = {Contrastive Distribution Matching for Amortized Sequential Monte Carlo in Discrete Diffusion},
  author = {Kim, Jaihoon and Yoon, Taehoon and Phunyaphibarn, Prin and Kim, Seungjun and Mardani, Morteza and Sung, Minhyuk},
  journal = {arXiv preprint arXiv:2605.23346},
  year = {2026}
}