Research map / Image generation / Fast diffusion sampling
DisCo-Diff: Enhancing Continuous Diffusion Models with Discrete Latents
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 DisCo-Diff, which augments continuous diffusion models with a few learnable discrete latents inferred by an encoder (trained end-to-end with the denoiser via Gumbel-Softmax) and modeled post hoc by an autoregressive transformer, reporting SOTA ODE-sampler FIDs on ImageNet-64/128 and gains on molecular docking.
Problem
Continuous diffusion models must map a unimodal Gaussian prior to a complex, multimodal data distribution, producing highly curved generative ODEs that are hard to learn and slow to sample.
Approach
Augment the diffusion denoiser with a small number of learnable discrete latent variables (global, low-cardinality codebooks) inferred by a ViT encoder trained jointly with the denoiser, so each discrete-conditioned sub-distribution has simpler, less curved score and ODE trajectories; a small autoregressive transformer then models the latent prior for sampling.
Key results
- DisCo-Diff reduces class-conditional ImageNet-64 FID with ODE sampler from 2.36 (EDM) to 1.65, and ImageNet-128 from 2.29 (VDM++) to 1.98 with ODE sampler.
- On PDBBind, DisCo-DiffDock-S raises full-set top-1 success from 32.9% to 35.4% and unseen-protein success from 13.9% to 18.5%.
- Unconditional ImageNet-64 FID improves from 3.94 (SCDM) to 3.70, and with Restart sampler ImageNet-64 FID reaches 1.22.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add learnable discrete latents to the diffusion model | EDM · ImageNet-64, class-conditioned, ODE sampler | FID: 2.36 to 1.65 | Helped |
| Add learnable discrete latents to the diffusion model | EDM · ImageNet-64, unconditional | FID: 3.94 -> 3.70 | Helped |
| Add learnable discrete latents to the diffusion model | VDM++ · ImageNet-128, class-conditioned, ODE sampler | FID: 2.29 to 1.98 | Helped |
| Add learnable discrete latents to the diffusion model | DiffDock-S · PDBBind, full test set | success rate: 32.9% to 35.4% | Helped |
| Add learnable discrete latents to the diffusion model | DiffDock-S · PDBBind, unseen proteins | success rate: 13.9% to 18.5% | Helped |
| Use DisCo-Diff with the VDM++ DDPM sampler | VDM++ · ImageNet-128, class-conditioned, DDPM sampler | FID: 1.88 to 2.80 | Hurt |
| Substitute VDM++ DDPM trajectories at smaller times | DisCo-Diff on ImageNet-128 · ImageNet-128, class-conditioned | FID: 2.80 to 1.73 | Helped |
| Replace discrete latents with continuous latents and use KLD weight 0.1 (ablation) | DisCo-Diff · ImageNet-64, class-conditioned, generative latent prior | FID: 1.65 -> 11.12 | Hurt |
| Replace discrete latents with continuous latents and use KLD weight 1 (ablation) | DisCo-Diff · ImageNet-64, class-conditioned, generative latent prior | FID: 1.65 -> 2.36 | Hurt |
| Set discrete-latent classifier-free guidance scale to 0 instead of 1 (ablation) | DisCo-Diff · ImageNet-64, class-conditioned, generative latent prior | FID: 1.65 -> 1.81 | Hurt |
| Increase discrete-latent classifier-free guidance scale from 1 to 2 (ablation) | DisCo-Diff · ImageNet-64, class-conditioned, generative latent prior | FID: 1.65 -> 2.33 | Hurt |
Surprising findings
- On ImageNet-128 with a DDPM stochastic sampler, DisCo-Diff (FID 2.80) is worse than VDM++ (1.88), which the authors attribute to discrete latents overemphasizing global information at small times.
- Continuous latents with a low KL weight (0.1) give a generative FID of 11.12 versus oracle 1.67, while discrete latents show only a minor generative-vs-oracle gap.
Limitations
- Two-stage training: encoder/denoiser are trained jointly, then the autoregressive prior is trained post hoc.
- Discrete latents hurt stochastic DDPM sampling on ImageNet-128, with the authors hypothesizing over-emphasis on global information.
- Hyperparameters (number of latents, codebook size) were not tuned due to computational constraints; evaluated mainly on standard benchmarks rather than text-to-image.
Experimental setup
Datasets: ImageNet PDBBind 2D mixture of Gaussians (synthetic)
Benchmarks: ImageNet-64 ImageNet-128 PDBBind
Metrics: FID NFE top-1 success rate within 2 Å RMSD runtime
Baselines: EDM PFGM++ VDM++ SCDM RIN CDM DiffDock-S DiffDock GNINA TankBind
Builds on: EDM VDM++ DiffDock PFGM++ Gumbel-Softmax
Closest papers by meaning
- Compositional Discrete Latent Code for High Fidelity, Productive Diffusion Models (NeurIPS 2025 · similarity 0.90)
- Unlocking Dataset Distillation with Diffusion Models (NeurIPS 2025 · similarity 0.90)
- Neural Diffusion Models (ICML 2024 · similarity 0.90)
- Plug-and-Play Diffusion Distillation (CVPR 2024 · similarity 0.90)
- Learning Diffusion Models with Flexible Representation Guidance (NeurIPS 2025 · similarity 0.90)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · 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.