Source-linked AI summary
Distillation of Discrete Diffusion through Dimensional Correlations
Satoshi Hayakawa, Yuhta Takida, Masaaki Imaizumi, Hiromi Wakaki, Yuki Mitsufuji
TL;DR
Discrete diffusion models commonly ignore dependencies between dimensions to keep high-dimensional joint distributions tractable, limiting few-step sampling. Di4C combines correlation-capable mixture models with distillation losses; experiments report improved or maintained generation quality across image and language tasks, alongside faster ImageNet sampling.
Problem
Discrete diffusion denoisers commonly factor across dimensions because high-dimensional joint distributions are intractable, but ignoring dependencies creates approximation errors when reducing sampling steps.
Source: §1 · p. 1, §2.2 · p. 3Method
Di4C uses mixture models to represent dimensional correlations and loss functions to distill many-step product-model denoising into fewer steps.
Source: §1 · p. 2Results
Takeaways & Limitations
Many-step product models can recover data distributions, and Di4C losses can distill them into few-step models when the student has enough expressive power.
Source: §7 · p. 9Although one-step distillation is proved in theory, empirical results show only about 2x fewer steps, leaving architecture and training-hyperparameter optimization as open needs.
Source: §7 · p. 9
Abstract
from arXiv · showhide
Diffusion models have demonstrated exceptional performances in various fields of generative modeling, but suffer from slow sampling speed due to their iterative nature. While this issue is being addressed in continuous domains, discrete diffusion models face unique challenges, particularly in capturing dependencies between elements (e.g., pixel relationships in image, sequential dependencies in language) mainly due to the computational cost of processing high-dimensional joint distributions. In this paper, (i) we propose "mixture" models for discrete diffusion that are capable of treating dimensional correlations while remaining scalable, and (ii) we provide a set of loss functions for distilling the iterations of existing models. Two primary theoretical insights underpin our approach: First, conventional models with element-wise independence can well approximate the data distribution, but essentially require {\it many sampling steps}. Second, our loss functions enable the mixture models to distill such many-step conventional models into just a few steps by learning the dimensional correlations. Our experimental results show the effectiveness of the proposed method in distilling pretrained discrete diffusion models across image and language domains. The code used in the paper is available at https://github.com/sony/di4c .
1. Introduction
Discrete diffusion denoisers commonly factor across dimensions for scalability, but this can miss dependencies needed for accurate few-step generation. Di4C combines a correlation-capable mixture model with distillation losses, supported by theoretical guarantees and experiments across image and language tasks.
1. Introduction
- Discrete diffusion uses product denoisers because high-dimensional joint distributions are intractable, but ignoring element-wise dependencies causes approximation errors when reducing sampling steps.The paper contrasts this challenge with continuous-domain approaches to faster sampling, including solvers and distillation.Source: §1 · p. 1
- Di4C uses a mixture model to represent dimensional correlations and loss functions to distill many-step product-model denoising into fewer steps.The paper states that these losses upper-bound the distance between an N-step teacher and a one-step student.Source: §1 · p. 2
- N-step sampling with product models approximates data distributions with O(1/N) total variation error, and this bound cannot be improved in a simple two-dimensional example.The result explains why many-step product models can work while motivating correlation modeling to reduce steps.Source: §1 · p. 2
- Di4C improves few-step sample-quality metrics on CIFAR-10, achieves a 2x ImageNet speed-up at comparable sample quality, and distills OpenWebText models while maintaining diversity.The OpenWebText result further distills an already-distilled masked diffusion model.Source: §1 · p. 2
2. Preliminaries
Discrete diffusion models generate data by reversing a noise-adding Markov process over a finite product space. For scalability, common denoisers factor across dimensions, limiting their ability to model dependencies in few-step generation.
2.1. Discrete diffusion models
- A discrete diffusion model uses a Markov process whose terminal distribution is tractable, then samples approximately from the data-conditioned reverse distribution.The state space is finite, and the data distribution is denoted q0.Source: §2.1 · pp. 2–3
- For a large product space, the forward transition is factorized over dimensions to support scalability.States consist of D entries drawn from a finite set S.Source: §2.1 · p. 3
2.2. Dimensional correlations in discrete diffusion
- Conventional denoisers use product models, making each dimension conditionally independent and reducing output length at the cost of representational capacity.This limitation is acute for one-step masked diffusion, where a fully masked input cannot yield a complex joint distribution from a dimensionally independent output.Source: §2.2 · p. 3
- Some continuous-time score-based losses use only marginal transition distributions, so learning backward transition rates does not capture dimensional correlations, even with a capable model.The passage identifies loss design, in addition to product-model parameterization, as a potential obstacle to learning correlations.Source: §2.2 · p. 3
3. Di4C for distilling discrete diffusion models
Di4C uses the composition of denoisers to capture dimensional correlations, then distills many-step discrete diffusion into fewer steps with mixture models and tailored losses.
3.1. Composition of diffusion denoisers for inducing dimensional correlations
- Composing dimensionally independent denoisers can induce dimensional correlations, which Di4C aims to compress into few-step sampling.The method distills compositions of well-trained denoisers into a student that learns those correlations.Source: §3.1 · p. 3
3.2. Mixture models to capture dimensional correlations
- A mixture over product models can represent any distribution on the product space while adding almost no inference overhead over a conventional product model.Each component is dimensionally independent; the mixture is indexed by a random variable λ.Source: §3.2 · p. 4
3.3. Consistency for distilling dimensional correlations
- Di4C combines a distillation loss that matches the student to the teacher at a small time δ with a consistency loss that learns correlations in teacher denoiser compositions.Reference distributions can come from data or from multiple teacher denoising steps.Source: §3.3 · p. 4, Figure 2 · p. 4
4. Theoretical analysis
The analysis bounds the approximation ability of many-step product models and establishes conditions under which Di4C losses distill a teacher into a few-step student.
4.1. Product models with multi-step sampling can approximate data distribution
- Product-model sampling has an O(1/N) total-variation upper bound and an Ω(1/N) lower bound in a two-dimensional example, so uniform error ε requires Ω(1/ε) steps.The theorem assumes correctly modeled marginals and regularity conditions.Source: §4 · p. 4, §4.1 · p. 5
4.2. Our losses can distill multi-step denoising models
- The proposed losses bound the total-variation distance between an N-step teacher and a one-step student in terms of the loss magnitudes.The guarantee applies when the student has enough expressive power; the bound vanishes if it perfectly learns the teacher composition.Source: §4.2 · pp. 5–6, §4 · p. 5
- Using data-generated q_t as the reference for times t < T avoids the cost of generating those reference samples with multi-step teacher inference.The passage states this choice should not harm training quality when the teacher is well-trained.Source: §4.2 · p. 6
5. Experimental results
Di4C was evaluated across pixel-space diffusion, masked image generation, and masked language modeling, with results indicating faster sampling and little mixture-model latency overhead.
5. Experimental results
- The experiments covered pixel-space discretized Gaussian diffusion on CIFAR-10, VQ-space masked image modeling on ImageNet, and masked diffusion language models on OpenWebText.They used different teacher models and state spaces across these settings.Source: §5 · p. 6, §5.1 · p. 6, §5.2 · p. 7
5.1. Discretized Gaussian diffusion on pixel space
- On CIFAR-10, the student substantially improved metrics over the teacher at 10 steps, while gains diminished as steps increased; the hybrid was best at 20 steps and matched teacher 40-step FID.At 20 total NFEs, adding predictor-corrector steps improved on vanilla denoising for both teacher and student, but allocating the budget to denoising performed better.Source: §5.1 · p. 6, Table 1 · p. 6, Table 2 · p. 6
5.2. Masked image modeling on VQ space
- 4-step Di4C closely matched 8-step teacher FID-IS performance on ImageNet 256 × 256, an approximate 2× sampling speed-up.The comparison varies classifier-free guidance coefficients.Source: §5.2 · p. 7, Figure 3 · p. 7
5.3. Masked diffusion language models
- Two Di4C rounds achieved more than 2× speed-ups while matching the teacher’s 64- or 128-step unconditional-generation performance, and Di4C models had higher diversity with comparable or better perplexity in conditional generation.One Di4C round gave a similar level of improvement to one SDTT round.Source: §5.3 · p. 8, Figure 4 · p. 8
5.4. Latency overhead of mixture modeling
- Mixture-model overhead was up to 5% on CIFAR-10 and OpenWebText, while the ImageNet student-teacher latency difference was within statistical noise.The authors attribute the ImageNet result likely to sampling, rather than logit inference, being the computational bottleneck.Source: §5.4 · p. 8, Table 3 · p. 8
6. Related work
Related work addresses faster sampling in continuous and discrete diffusion, as well as methods that incorporate dimensional correlations. The paper positions Di4C as compatible with some discrete-diffusion approaches while noting opportunities to combine its losses with other correlation models.
Speeding up continuous diffusion models
- Continuous-diffusion distillation and consistency methods enable single- or few-step sampling, mostly for probability-flow ODEs; Gaussian mixtures have also reduced steps in non-ODE diffusion (Li et al., 2024).The cited mixture work addresses limited expressive power of unimodal Gaussian denoisers.Source: §6 · pp. 8–9
Speeding up discrete diffusion models
- Discrete-diffusion speedups include optimized sampling schedules, SDTT’s approximately 32× speed-up over its teacher, and MaskGIT confidence-based sampling; Di4C was demonstrated with the latter two approaches.Source: §6 · p. 9
Related work on dimensional correlations
- Other correlation methods combine autoregressive and masked-diffusion models or use energy-based sampling, but can add time or memory overhead; Di4C uses simpler sampling, and combining its model-agnostic losses with those methods remains future work.The overhead is attributed to an additional non-diffusion model and, for the energy-based approach, importance sampling.Source: §6 · p. 9
7. Conclusion
The conclusion presents Di4C as a way to distill pretrained discrete diffusion models using losses and models that capture dimensional correlations. It reports cross-domain results while identifying the gap between theoretical one-step distillation and current empirical speed-ups.
Conclusion
- The paper states that many-step product models can recover a data distribution and that Di4C losses can distill them into few-step models when the student has enough expressive power.Source: §7 · p. 9
- Experiments improved 10-step CIFAR-10 sample quality, achieved a 2× speed-up on ImageNet 256 × 256, and distilled already-distilled OpenWebText models while maintaining generation diversity.Source: §7 · p. 9
- Although one-step distillation is proved in theory, empirical results achieve only about 2× fewer steps so far; the authors point to λ-related architecture and training-hyperparameter optimization as needed work.Iterated Di4C training is identified as a promising candidate.Source: §7 · p. 9
Impact Statement
The impact statement acknowledges potential social impacts associated with generative AI but reports no specific problems particularly caused by this mostly theoretical work.
Impact Statement
- The authors identify no specific problems particularly caused by this mostly theoretical work.Source: Impact Statement · p. 9
- The statement situates potential social impacts in the broader generative-AI field rather than identifying a paper-specific impact.Source: Impact Statement · p. 9
- The impact assessment characterizes the work as mostly theoretical.Source: Impact Statement · p. 9
A. Training techniques for Di4C
The appendix develops practical loss-computation and training techniques for Di4C, then reviews the distributional inequalities and Markov-chain equations used in its analysis.
A. Training techniques for Di4C
- Monte Carlo integration and control variates make Di4C loss computation scalable.The control variate is chosen to approximate the target while keeping its expectation tractable.Source: A. Training techniques for Di4C · p. 13, §A.2 · p. 14
A.3. Auxiliary losses
- The datapoint loss computes data-distribution negative log-likelihood when data are available, while the marginal loss is easier to compute under a teacher-marginal accuracy assumption.Mixing the marginal and correlation losses with time-dependent α_t was more efficient than either pure marginal or pure data loss.Source: A. Training techniques for Di4C · p. 13, §A.3 · p. 15
A.1. Surrogate of distillation loss
- A convex upper bound on the distillation loss is a practical surrogate, and it vanishes when the student matches the teacher over [0, δ].The authors note this is feasible when the student has enough expressive power.Source: §A.1 · p. 13
A.2. Surrogate of consistency loss
- A dimensionally independent product model serves as a control variate for consistency-loss Monte Carlo estimates, with the residual capturing dimensional correlations.The decomposition uses one term to learn correlations and another to stabilize the approximation using more information about the target distribution.Source: §A.2 · p. 14
A.4. Iterated Di4C training
- Di4C losses can use mixture teachers, enabling multiple distillation rounds; the marginal loss is the exception and instead requires replacing q0|t in its estimator.The authors relate this generalization to multi-round SDTT (Deschenaux & Gulcehre, 2025).Source: §A.4 · p. 15
B. Kullback–Leibler divergence and total variation distance
- The appendix reviews KL divergence and total variation, including Pinsker’s inequality and the convexity of KL divergence used in the analysis.It also states a triangle-like inequality for total variation between composed distributions.Source: B. Kullback–Leibler divergence and… · p. 15
C. Continuous-time Markov chains and Kolmogorov equations
- For finite-state continuous-time Markov chains, short-time transition rates provide the starting point for deriving forward and backward Kolmogorov equations.The appendix tracks continuity and nonzero assumptions in reorganizing earlier derivations.Source: C. Continuous-time Markov chains and… · p. 15, §C.1 · p. 15
C.1. Kolmogorov equations in the general case
- Continuity of the transition rates and conditional distributions yields the forward and backward Kolmogorov equations on the interior time interval.The appendix notes the sums are finite because the state space is finite.Source: §C.1 · p. 16
C.2. Kolmogorov equations for factorized forward processes
This section analyzes factorized forward processes and the product denoisers built from their marginals. It establishes a 1/N convergence rate with a matching lower-bound example, then supplies supporting proof arguments.
D.3. Proof of Proposition 5
- For positive-probability states, the reverse transition rate is obtained from the forward rate weighted by the ratio of state probabilities; it can be nonzero only when states differ in at most one entry.The probability ratio is described as a discrete counterpart of the score function.Source: §C.2 · p. 17
C.2. Kolmogorov equations for factorized forward processes
- Under the stated assumptions, product denoisers with correct marginals have a 1/N convergence rate, and a two-dimensional binary example gives a matching lower bound.The stated rate is for analytical sampling with timesteps ti = iT/N.Source: §C.2 · p. 17
C.3. Lower bound of Theorem 3
- In the uniform-diffusion example, the authors establish an Ω(1/N) error between the analytical and true denoisers.The example uses two binary dimensions and satisfies Assumptions A and B.Source: §C.3 · pp. 17–18
D.1. Proof of Theorem 2
- The proof of Theorem 2 combines consistency-loss comparisons across adjacent steps with Pinsker’s inequality and the triangle inequality.The proof uses KL convexity in the consistency comparison.Source: §D.1 · p. 18
D.2. Proof of Proposition 4
- The proof of Proposition 4 handles equal conditional distributions and equal marginals separately, then combines the bounds using the total-variation triangle inequality.Source: §D.2 · p. 19
D.4. Proof of first half of Theorem 3
The proof bounds the global error by controlling local reverse-transition error and composing the resulting steps. The appendix also describes the symmetric binary example and explains the control-variate mechanism.
D.4. Proof of first half of Theorem 3
- A local reverse-transition approximation bound and a composition inequality yield the first-half Theorem 3 estimate.The proof bounds each step and sums the resulting terms across the composed sampler.Source: §D.4 · pp. 20–22
D.5. Proof of Proposition 6
- In the symmetric binary example, the probabilities of aa recover the information in the distributions considered, simplifying the lower-bound calculations.The proof computes analytical-sampler probabilities and compares their recurrence with the true distribution’s recurrence.Source: §D.5 · pp. 22–24
E. Control variates
- A control variate replaces direct Monte Carlo estimation with an estimate of the residual plus a tractable expectation, ideally reducing variance while remaining unbiased.Source: E. Control variates · p. 27
E.1. Marginal-matching product model as control variate
- A product model with matching marginals approximates the mixture model, leaving the difference to represent dimensional correlations for Monte Carlo estimation.The product model’s expectation is easy to compute or approximate.Source: §E.1 · p. 27
E.2. Derivations of dimension-wise computable control variates for mixture model
The appendix derives dimension-wise computable control variates for mixture-model losses, explains why product-model cross-entropy depends only on marginals, and describes the sampling schemes used in experiments.
E.2. Derivations of dimension-wise computable control variates for mixture model
- A Jensen-based upper bound and a dimension-wise computable alternative support a decomposition of mixture-model loss computation into Monte Carlo and tractable terms.The decomposition targets high variance when sampling only from q.Source: §E.2 · pp. 27–28
E.3. Product teacher model as control variate
- For a product distribution, cross-entropy is the same for two distributions with identical marginals, so it can be computed using those marginals alone.Source: §E.3 · p. 28
F.1.1. SAMPLING SCHEMES
- The experiments use τ-leaping and analytical sampling; the authors report that τ-leaping deteriorates at fewer steps, while analytical sampling suffers less from discretization.Both schemes are dimensionally independent.Source: §F.1.1 · p. 29
F.1.2. ADDITIONAL EXPERIMENTAL RESULTS
This section presents additional experiments and implementation details, covering sampling efficiency and behavior, masked-diffusion experiments and ablations, generated-sample comparisons, and SDTT checkpoint distillations with and without data loss.
F.1.2. ADDITIONAL EXPERIMENTAL RESULTS
- The hybrid model beats other models in 20-step FID and matches the teacher’s 40-step FID, while the student alone worsens at 40 steps.The authors hypothesize that increasing dimensional independence in the true denoiser helps explain the step-dependent results, and note that other forward diffusions or noise schedules warrant investigation.Source: §F.1.2 · pp. 29–30
- FID results show approximately 1.4× acceleration at 10–20 steps, but the method performs poorly at 2–4 steps.These comparisons use 10,000 generated samples, fewer than in the main-body evaluation.Source: §F.1.2 · p. 29, Table 4 · p. 30
- 40-step analytical sampling achieves better FID than 1000-step τ-leaping with the teacher model.The appendix compares these two sampling strategies and reports the complete evaluation in Table 5.Source: §F.1.2 · p. 29, Table 5 · p. 30
F.1.3. IMPLEMENTATION AND TRAINING
- The CIFAR-10 models use a U-net denoiser and dimension-wise analytical sampling; the student is finetuned from the teacher checkpoint with an added λ embedding.The appendix also describes CIFAR-10 training settings and shows generated examples corresponding to Table 1 in Figure 5.Source: §F.1.3 · pp. 30–31, §F.1.5 · p. 32, Figure 5 · p. 33
F.1.4. ABLATION STUDY
- Control variates consistently help, non-constant α_t functions work better than α_t = 0 or 1, and λ batch size trades off FID against IS.A larger λ batch tends to improve FID, whereas a smaller one tends to improve IS in 10- or 20-step sampling.Source: Table 6 · p. 31, §F.1.4 · p. 31, Table 7 · p. 32
F.2.1. MASKED DIFFUSION MODELING
- Masked diffusion uses independently masked tokens, while confidence-based sampling unmasks tokens by ranked confidence and can condition the product model on a sampled λ.The ImageNet setup uses an arccos masking schedule; discrete classifier-free guidance increases over the sampling process.Source: §F.2.1 · p. 32, §F.2.2 · pp. 33–34, §F.2.3 · p. 34
F.2.4. IMPLEMENTATION AND TRAINING
- The ImageNet appendix describes a transformer mixture-model adaptation and reports selected CFG coefficients, FID/IS, and Precision/Recall alongside generated samples.The teacher uses a 24-layer transformer; the mixture model adds a λ embedding, and Figure 6 shows samples for Ostrich, Burger, and Volcano labels.Source: §F.2.4 · p. 34, §F.2.5 · p. 35, Table 8 · p. 35
F.3.1. DIFFUSION MODELING AND SAMPLING
- The language-modeling appendix uses D = 1024 tokens with a GPT-2 codebook and describes analytical sampling for masked diffusion; its teacher is a 169M-parameter transformer with digitized timesteps.The teacher checkpoints are from Deschenaux & Gulcehre (2025), and the timestep spacing is ∆t = 1/1024.Source: §F.3.1 · p. 35, §F.3.2 · p. 37
F.3.2. IMPLEMENTATION AND TRAINING
- Training mostly followed Deschenaux & Gulcehre (2025), using Adam at 3 × 10^-5, EMA weight decay 0.9999, and a 500-iteration linear warm-up.Each Di4C round ran for 100K iterations on 2 A6000 GPUs, with minibatch size 2 and λ-batch size 16.Source: §F.3.2 · p. 37
- Conditional-generation Self-BLEU used five 100-token continuations per prompt, 256 WebText prompts, and 4-gram BLEU with equal weighting and method1 smoothing.Prompts were the first 50 tokens of WebText datapoints; inputs used the GPT-2 tokenizer, and BLEU computation followed Zhu et al. (2018) and Bird et al. (2009).Source: §F.3.2 · p. 37
- Figures 7 and 8 compare Di4C distillations of SDTT checkpoints with and without data loss.The captions specify the comparison but provide no outcome values.Source: Figure 7 · p. 38, Figure 8 · p. 38
F.3.3. ADDITIONAL EXPERIMENTAL RESULTS
Additional experiments find that adding data loss leaves results almost unchanged, while MAUVE shows no significant performance decay from the teacher. The section also presents paired qualitative samples of SDTT-7 and its Di4C2 distillation.
F.3.3. ADDITIONAL EXPERIMENTAL RESULTS
- In the data-loss ablation, SDTT-6/7 + Di4C-d using loss (72) with αt = 0.1g(t) performed almost the same as no-data-loss variants, whose generative perplexities were slightly better.In Figure 7, the SDTT-7 + Di4C curve is hidden behind the SDTT-7 + Di4C-d curve.Source: §F.3.3 · p. 38
- The MAUVE evaluation shows no significant performance decay from the teacher model.The setting matches unconditional generation in Section 5.3.Source: §F.3.3 · p. 38
F.3.4. GENERATED SAMPLES
- The generated-samples subsection pairs SDTT-7 and SDTT-7 + Di4C2 outputs conditioned on the first 50 tokens of the abstract.The displayed examples include 4-step outputs from both models.Source: §F.3.4 · pp. 38–39