Neural-CTMC

Neural CTMC: Discrete Diffusion via Decoupled Jump Timing and Direction

This repository contains the inference checkpoint and demo code for Neural CTMC, a discrete diffusion model based on continuous-time Markov chains (CTMCs). Unlike prior methods that parameterize the reverse rate matrix as a monolithic object, Neural CTMC separately parameterizes the exit rate (when to jump) and the jump distribution (where to jump) via two dedicated network heads, aligning the parameterization with the intrinsic CTMC decomposition.

This checkpoint is trained on OpenWebText with a uniform forward process and is, to our knowledge, the first open-source checkpoint for a uniform-noise discrete diffusion language model.

Model Details

Performance

Generative perplexity (scored by Gemma-2, lower is better) on OpenWebText:

Neural CTMC achieves the best generative perplexity among equal-budget (262B) methods across all step counts, and remains competitive with SEDD despite using 2.6x fewer training tokens.

Usage

Requirements

pip install torch transformers

Quick Start

from demo_infer import CTMCHFModel

model = CTMCHFModel.from_pretrained(
    "owt_uniform.pt",
    device="cuda",
    tokenizer_name="gpt2",
)

texts = model.generate(
    n_samples=3,   # number of samples to generate
    n_steps=128,   # Euler discretization steps
    T=1.0,         # diffusion time horizon
)

for i, text in enumerate(texts):
    print(f"[Sample {i+1}]")
    print(text)

Command Line

# Generate 5 samples with 128 Euler steps on GPU 0
GPU=0 bash run.sh

You can also call the script directly:

python demo_infer.py \
    --checkpoint owt_uniform.pt \
    --n_samples 5 \
    --n_steps 128 \
    --T 1.0 \
    --device cuda \
    --output output/samples.txt

How It Works

The model generates text through reverse diffusion over discrete token sequences using the Euler sampler:

1. Initialize a sequence of uniformly random tokens of length 512.
2. Iteratively denoise for n_steps Euler steps: at each step, the model predicts per-token exit rates $\lambda^\theta_t$ and a jump distribution $r^\theta_t$ over the vocabulary, then stochastically updates tokens via the CTMC reverse process.
3. Decode the final token sequence using the GPT-2 tokenizer.

The key insight is that the ELBO decomposes into a Poisson KL for jump timing and a categorical KL for jump direction, enabling the model to learn these two aspects with separate heads.

File Structure

.
├── README.md            # This file
├── owt_uniform.pt       # Model checkpoint (~969 MB)
├── demo_infer.py        # Inference script with CTMCHFModel class
└── run.sh               # Convenience launch script

Citation

If you find this model useful, please cite our work:

@article{li2025neuralctmc,
  title={Neural Continuous-Time Markov Chain: Discrete Diffusion via Decoupled Jump Timing and Direction},
  author={Jingyuan Li and Xiaoyi Jiang and Fukang Wen and Wei Liu and Renqian Luo and Yi Zhu and Zuoqiang Shi and Pipi Hu},
  year={2025}
}

License

This project is licensed under the MIT License.