PaperlayerView exact revision on arXiv

Source-linked AI summary

Score-based Continuous-time Discrete Diffusion Models

Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, Hanjun Dai

arXiv:2211.16750v2cs.LG

TL;DR

Score-based diffusion is difficult to extend to categorical data because the continuous-space score is undefined there. The paper introduces a continuous-time Markov-chain diffusion with categorical score matching and analytical reverse sampling, achieving promising results across synthetic, music, and image benchmarks.

  • Problem

    Score-based diffusion requires a score function that is not defined for discrete variables, while prior discrete extensions rely on ELBO approximations and fixed sampling strategies.

  • Method

    The paper models categorical diffusion with a stochastic jump process, learns reverse dynamics through categorical ratio matching, and develops analytical reverse sampling based on conditional marginals.

  • Results

    The proposed SDDM achieves promising results on synthetic and real-world music and image benchmarks, approaching VQ-based categorical diffusion and improving over D3PM and τLDR in the same VQ space.

  • Takeaways & Limitations

    Continuous-time score-based diffusion can be formulated for categorical spaces with tractable learning and reverse-sampling procedures.

  • Takeaways & Limitations

    Conditional-marginal parameterization requires trade-offs between computational cost and architectural flexibility, and the energy-based parameterization can require O(D × C) network evaluations.

Abstract

from arXiv · show

Score-based modeling through stochastic differential equations (SDEs) has provided a new perspective on diffusion models, and demonstrated superior performance on continuous data. However, the gradient of the log-likelihood function, i.e., the score function, is not properly defined for discrete spaces. This makes it non-trivial to adapt \textcolor{\cdiff}{the score-based modeling} to categorical data. In this paper, we extend diffusion models to discrete variables by introducing a stochastic jump process where the reverse process denoises via a continuous-time Markov chain. This formulation admits an analytical simulation during backward sampling. To learn the reverse process, we extend score matching to general categorical data and show that an unbiased estimator can be obtained via simple matching of the conditional marginal distributions. We demonstrate the effectiveness of the proposed method on a set of synthetic and real-world music and image benchmarks.

1 INTRODUCTION

Discrete diffusion extensions have promising results but remain tied to finite corruption steps and fixed reverse sampling, while continuous-time score-based modeling is difficult because scores are undefined for discrete variables. This paper addresses both issues with a continuous-time Markov-chain formulation, categorical score matching, and new sampling methods.

  • Discrete diffusion extensions use alternative corruption operations but retain finite corruption and restoration sequences with fixed reverse sampling strategies.
  • Continuous-time score-based modeling is non-trivial for discrete variables because the log-likelihood gradient is undefined.
  • The paper extends score functions to generic categorical variables and formulates continuous-time discrete diffusion with a continuous-time Markov chain.
  • It introduces categorical ratio matching, analytical reverse sampling based on conditional marginals, architectural choices including a hollow neural model, and evaluations on synthetic, music, and image benchmarks.

2 BACKGROUND

Diffusion models connect a forward corruption process with a reverse generative process, and continuous-time formulations use SDEs indexed over a continuous time interval. The standard score-based SDE does not extend directly to discrete variables because the discrete score is undefined.

  • Diffusion models transform data into a reference distribution through a forward Markov process and recover the data distribution with a backward Markov process.
  • Discrete forward corruption kernels can be uniform, discrete Gaussian, or another choice, while continuous models commonly use Gaussian noise.
  • The backward kernel is generally intractable and is typically parameterized by a neural network and learned with ELBO or score matching.
  • Continuous-time diffusion indexes t over [0, T] and defines forward and reverse dynamics through SDEs and the score function.
  • The score-based SDE does not characterize discrete diffusion because ∇x log pt(x) is undefined for a discrete variable.

3 CONTINUOUS TIME DISCRETE SCORE MATCHING

The paper replaces continuous-space score gradients with a continuous-time Markov-chain framework whose reverse rates depend on probability ratios. It learns those ratios by matching singleton conditional distributions, yielding a tractable categorical score-matching objective.

  • 3.1 CONTINUOUS TIME MODELING: The forward process is a continuous-time Markov chain on a finite categorical space, with transitions characterized by time-dependent rate matrices.
  • 3.1 CONTINUOUS TIME MODELING: Choosing rate matrices to approach a tractable reference distribution makes the reversed Markov process a generative path toward the data distribution.
  • 3.1 CONTINUOUS TIME MODELING: The reverse process is determined by reverse rate matrices, whose closed-form expression requires the marginal probability ratio qt(y)/qt(x).
  • 3.2 CATEGORICAL RATIO MATCHING: Because this ratio is generally intractable, the method generalizes binary ratio matching to categorical variables through singleton conditional distributions.
  • 3.2 CATEGORICAL RATIO MATCHING: Matching singleton conditionals is sufficient to match the joint distribution, and a time-dependent neural network learns these conditionals along the forward process.
  • 3.2 CATEGORICAL RATIO MATCHING: The categorical ratio-matching loss has a tractable simplified form that can be efficiently estimated with Monte Carlo sampling.

4 CONTINUOUS TIME DISCRETE SAMPLING

The paper develops reverse sampling for continuous-time discrete diffusion, using parallel Euler updates and an analytical alternative that avoids Euler simulation error. The approach addresses the reverse process’s non-factorization while retaining efficient sampling.

  • Discrete-time reverse sampling: Reverse jump rates depend on the current values of other dimensions, so exact dimension-wise parallel simulation is unavailable.This dependence makes the reverse process non-factorizable even when the learned conditional distribution is known.
  • Discrete-time reverse sampling: Parallel Euler updates sample every dimension simultaneously, making reverse-process simulation more efficient.The method fixes the reverse rate over a short interval, clips transition quantities to valid probabilities, and samples each dimension independently.
  • Analytical sampling: Large Euler steps can shift samples away from the correct time-slice marginal because the reverse rate is held fixed across each interval.The paper notes that standard correctors could reduce this approximation error but would increase computational cost.
  • Analytical sampling: Analytical sampling replaces Euler’s explicit approximation with an implicit parameterization based on conditional marginals.Because the forward conditional transition is tractable, the resulting transformation can be computed efficiently once the conditional distribution is known.
  • Analytical sampling: The resulting analytical reverse sampler avoids the simulation error introduced by Euler’s method.The paper trains the implicit conditional model with categorical ratio matching before using the analytical reverse-process expression.

5 PARAMETERIZATION

The parameterization section compares energy-based, masked, and hollow-Transformer architectures for enforcing the required exclusion of a coordinate’s current value. These designs trade flexibility and computational cost, with the hollow Transformer requiring only O(1) feed-forward evaluations.

  • Design constraint: Predictions for dimension d must not depend on its current value x_d^t, although they may use the other coordinates.This constraint prevents information leakage that would make the categorical ratio-matching objective trivial.
  • Energy-based models: Energy-based models offer the most flexible parameterization by assigning energies with an arbitrary neural network.Their conditional marginals require O(D × C) evaluations of the network, which is computationally prohibitive for high-dimensional data.
  • Masked models: Masked models preserve flexible neural networks while replacing one coordinate with a MASK token and reducing evaluation cost to O(D).Reducing the number of dimensions can further lower evaluations at the cost of increasing vocabulary size C.
  • Hollow Transformer: The proposed hollow Transformer requires only O(1) feed-forward evaluations while enforcing zero diagonal entries in the conditional-probability Jacobian.This removes the dimensionality- and vocabulary-dependent evaluation scaling of the other parameterizations.
  • Hollow Transformer: Autoregressive and hollow masking are cited as techniques for enforcing the required dependency structure.Autoregressive masking creates a triangular Jacobian, whereas hollow masking permits full-context interaction through a single dense layer.

6 EXPERIMENTS

Experiments evaluate the proposed diffusion approach on synthetic binary data, CIFAR10 images, and monophonic music. Results indicate strong sample quality, flexible reverse sampling, and robustness with fewer sampling steps.

  • Synthetic data: On synthetic 32-dimensional binary distributions, categorical ratio matching learns EBMs with consistently lower exponential Hamming MMD than the baselines.Samples decoded from Gray codes show distributions similar to the ground truth.
  • Synthetic data: Across three parameterizations, 3-layer transformers achieve comparable performance, while masked and hollow models offer better quality-speed trade-offs than EBMs.
  • Image modeling: On CIFAR10, categorical discrete modeling is harder than ordinal modeling, but the proposed approach approaches VQ-based categorical diffusion and improves on D3PM and τLDR in the same VQ space.
  • Image modeling: The analytical sampler achieves reasonable image quality in fewer reverse-sampling steps, while continuous-time modeling remains more robust than D3PM when using fewer steps.The forward Euler sampler may require a corrector for better performance at low step counts.
  • Monophonic music modeling: The music benchmark contains 6,000 training sequences and 973 evaluation sequences of length 256, with a scrambled vocabulary of 128 notes plus a rest.The scrambled vocabulary removes ordering information and creates a challenging categorical modeling task.

7 LIMITATION AND CONCLUSION

The paper presents score-based continuous-time diffusion for categorical spaces, with learning and sampling procedures designed for discrete variables. It reports promising results while identifying parameterization, ordinal-structure, and mixed-space limitations.

  • Conclusion: The paper introduces a learning and sampling paradigm for continuous-time diffusion models in categorical discrete spaces.The approach extends score matching to categorical variables and uses a continuous-time Markov-chain formulation.
  • Conclusion: Score matching aligns the reverse process with the posterior of the forward process for categorical discrete variables.
  • Conclusion: The method introduces sampling algorithms naturally through its score-based learning formulation.The paper contrasts its analytical sampling strategy with prior ELBO-based learning and sampling procedures.
  • Limitations: The main bottleneck is parameterizing conditional marginals while balancing computational cost against architectural flexibility.
  • Limitations: General categorical score matching does not exploit prior ordinal structure, and mixed-space data require a unified continuous-discrete score-matching treatment.

B.4 PROOF FOR PROPOSITION 3.4

The proof establishes a simplification of the score-matching objective using conditional distributions, yielding a tractable loss; it also records an empirically effective cosine-style noise schedule.

  • Proof: The conditional distribution on dimension d does not depend on the value x_d, enabling the subsequent objective simplification.
  • Proof: Substituting the conditional-distribution result into the original score-matching loss yields a simplified, tractable loss function.
  • Noise schedule: The empirically effective schedule β(t) provides cosine-style noise levels for the forward process.The stated motivation is to maintain a reasonable noise level that contributes to sample quality.

C.2 SYNTHETIC DATA

The synthetic-data section evaluates the method on Gray-code representations of 2D distributions and compares parameterization and sampling choices. The ablation reports reasonable results across combinations, with mixed sampler-model comparisons.

  • Synthetic data: Figure 4 visualizes the true data in two-dimensional space by decoding the Gray codes.
  • Ablation: All tested parameterization-and-sampler combinations achieve reasonable results in the synthetic-data ablation.
  • Ablation: Comparisons among different samplers and models are mixed, while masked or hollow parameterizations can be more efficient.

C.3 EXPERIMENTS ON CIFAR10

The CIFAR10 experiments model images through categorical discrete latent representations and evaluate the resulting generation system. The latent reconstruction reaches FID 9.05 and IS 9.67.

  • Vector quantization: A VQ-VAE maps images into categorical latent tokens with a vocabulary size of 512.The encoder produces spatially downsampled latent codes for image modeling.
  • Vector quantization: The VQ-VAE uses GAN and perceptual losses alongside the ELBO objective during training.
  • CIFAR10 results: FID 9.05 and IS 9.67 are achieved when reconstructing 32×32 CIFAR10 images from 8×8 latent codes.
  • Parameterization and training: The CIFAR10 diffusion model uses masked modeling with a BERT-base Transformer backbone and predicts logits for masked positions.
  • Parameterization and training: The CIFAR10 model is trained for 700k steps using TPU-v4 hardware and an exponentially averaged parameter set for evaluation.

C.4 MUSIC DATASET

The model uses a hollow Transformer with specified architectural dimensions and is trained on 2×2 TPU-v4 chips until convergence.

  • Model architecture: The proposed hollow Transformer uses 6 layers, embedding size 256, 8 attention heads, and hidden dimension 2048.These settings apply to each Transformer component.
  • Model architecture: Time embeddings are concatenated with the other tokens before feed-forward processing.
  • Training setup: Training on 2×2 TPU-v4 chips reached convergence in about 12 hours, or roughly 2 million steps with batch size 64.
Loading 2211.16750v2…