Research map / Image generation / Fast diffusion sampling
Jump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Jump Your Steps, which optimizes the timestep schedule of discrete diffusion model sampling by minimizing a KL-divergence upper bound (KLUB) on compounding decoding error from tau-leaping, with no extra inference cost.
Problem
Parallel (tau-leaping / k-Gillespie) sampling in discrete diffusion models is fast but introduces compounding decoding error because tokens are sampled independently, degrading quality at low NFE.
Approach
Derive a tractable KL-divergence upper bound (KLUB) on compounding decoding error using Girsanov's theorem for CTMCs, simplify it via maximizing coarse-to-fine KL and closed-form expected transition counts, and greedily search breakpoints with a hierarchical split and golden-section search.
Key results
- On CIFAR-10, JYS yields better FID than the uniform schedule at every NFE from 32 to 256 (FID computed on 50k samples).
- On text generation with SEDD, JYS gives better generative perplexity than uniform schedules at the same NFE from 16 to 256, using GPT-2-large as evaluator.
- On Countdown, JYS produces fewer rule violations than the uniform schedule at the same NFE.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the uniform schedule with the JYS schedule at the same NFE | SEDD on Countdown · Countdown | rule-violation rate | Helped |
| Replace the uniform schedule with the JYS schedule at the same NFE | pretrained discrete diffusion model · CIFAR-10 | FID | Helped |
| Replace the uniform schedule with the JYS schedule at the same NFE | pretrained discrete diffusion model · Lakh pianoroll | Hellinger distance to 512-NFE note distributions | Helped |
| Replace the uniform schedule with the JYS schedule at the same NFE | SEDD-small and SEDD-medium · OpenWebText text generation | generative perplexity | Helped |
| Vary the number of samples used for KLUB computation from 4 to 1024 (ablation) | JYS on CIFAR-10 · CIFAR-10 | sampling quality | No clear effect |
Surprising findings
- The KLUB estimate barely changes when the number of Monte Carlo samples drops from 1024 to 16, so schedule optimization can be done in under a minute with few samples.
Limitations
- KLUB is derived for state-independent CTMCs and only approximates the state-dependent reverse process of DDMs.
- Several approximations (e.g., single transition per dimension, coarse-to-fine proxy) are used, with the exact-ness only holding under assumptions such as the absorbing kernel.
Experimental setup
Datasets: Countdown (synthetic) CIFAR-10 Lakh Pianoroll OpenWebText
Benchmarks: CIFAR-10 Lakh Pianoroll OpenWebText generation Countdown
Metrics: FID Hellinger distance generative perplexity rule violation rate
Baselines: uniform schedule predictor-corrector sampler
Builds on: tau-leaping k-Gillespie Girsanov theorem SEDD Align Your Steps
Closest papers by meaning
- Align Your Steps: Optimizing Sampling Schedules in Diffusion Models (ICML 2024 · similarity 0.91)
- Learnable Sampler Distillation for Discrete Diffusion Models (NeurIPS 2025 · similarity 0.90)
- A Simple Early Exiting Framework for Accelerated Sampling in Diffusion Models (ICML 2024 · similarity 0.90)
- Fast Sampling via Discrete Non-Markov Diffusion Models with Predetermined Transition Time (NeurIPS 2024 · similarity 0.90)
- On the Error Correcting Effects of Stochasticity in Discrete Diffusion (NeurIPS 2026 · similarity 0.89)
- Fast Solvers for Discrete Diffusion Models: Theory and Applications of High-Order Algorithms (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.