PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion on the Probability Simplex

Griffin Floto, Thorsteinn Jonsson, Mihai Nica, Scott Sanner, Eric Zhengyu Zhu

arXiv:2309.02530v2cs.LGstat.ML

TL;DR

Diffusion models face a mismatch between continuous Gaussian noising and discrete categories. This paper maps an Ornstein–Uhlenbeck process onto the probability simplex and reports exact SDE solutions for simplex and unit-cube diffusion, while identifying further evaluation as future work.

  • Problem

    Gaussian noising in continuous spaces leaves the question of how to add noise when the underlying categories are discrete.

    Source: §1 · p. 1

  • Method

    The forward process maps an Ornstein–Uhlenbeck process onto the probability simplex via an additive logistic transformation, yielding an efficiently sampled logistic-Gaussian transition kernel.

    Source: §3.3 · p. 3

  • Results

    The authors report exact solutions for the simplex and unit-cube SDE dynamics, with the method fitting the common diffusion training paradigm.

    Source: §6 · p. 4

  • Takeaways & Limitations

    Probability-space states represent category probabilities, while the product construction keeps higher-dimensional diffusion within the unit cube with a closed-form score.

    Source: §1 · p. 1, §5.2 · p. 4

    The authors identify further testing on more complex datasets and evaluation of categorical-distribution properties as future work.

    Source: §6 · p. 4

Abstract

from arXiv · show

Diffusion models learn to reverse the progressive noising of a data distribution to create a generative model. However, the desired continuous nature of the noising process can be at odds with discrete data. To deal with this tension between continuous and discrete objects, we propose a method of performing diffusion on the probability simplex. Using the probability simplex naturally creates an interpretation where points correspond to categorical probability distributions. Our method uses the softmax function applied to an Ornstein-Unlenbeck Process, a well-known stochastic differential equation. We find that our methodology also naturally extends to include diffusion on the unit cube which has applications for bounded image generation.

1. Introduction

Diffusion models progressively add noise and learn a reverse process, but Gaussian noising in continuous spaces leaves a challenge for discrete categories. The paper addresses this by diffusing over probability vectors on the simplex, while noting advantages of diffusion models over other generative approaches.

1. Introduction

  • Diffusion models progressively add noise to data and learn a reverse stochastic differential equation to generate samples.The forward process transforms a complex data distribution into a simpler distribution that is easier to sample from.Source: §1 · p. 1
  • Compared with GANs, diffusion models offer an exact likelihood interpretation and do not require adversarial training.Source: §1 · p. 1
  • Diffusion models avoid non-overlapping data and generated distributions, while discretized normalizing flows face Jacobian-determinant computation restrictions (Chen et al., 2018).Source: §1 · p. 1
  • Mapping categories to probability vectors on the simplex makes discrete observations amenable to continuous diffusion.Each simplex point represents the probability of choosing a category.Source: §1 · p. 1

2. Background

Score-based diffusion describes a forward stochastic process that maps data to an easy-to-sample limiting distribution. Generation approximates the data score and uses the reverse SDE, trained with a score-matching objective.

2. Background

  • Typically, the parameter-free forward process maps data to an easy-to-sample limiting distribution, and its time reversal is also an SDE.Source: §2.1 · p. 1
  • Generation approximates the score ∇x log pt(xt) and uses the reverse SDE to sample from the generative model.Source: §2.1 · p. 2
  • The score-matching objective uses λ(t) as a weighting function and pst(xt|xs) as the transition kernel; time is commonly discretized into uniform steps.Source: §2.1 · p. 2

3. Method

The method constructs diffusion on simplex-valued categorical probabilities by mapping an Ornstein–Uhlenbeck process through a logistic transformation. It also specifies how discrete data are represented and relaxed for training.

3.1. The Logistic-Normal Distribution on the Probability Simplex

  • The logistic-normal is a distribution over simplex probabilities obtained by softmaxing Gaussian logits; in d=2, this reduces to a sigmoid mapping R to [0,1].Source: §3.1 · p. 2

3.2. The Ornstein-Unlenbeck Process

  • The OU process is real-valued, has a normal distribution at time t, and has a limiting distribution uniquely determined by θ.Source: §3.2 · p. 2

3.3. Diffusion on the Probability Simplex

  • Mapping an OU process through the additive logistic transformation yields a simplex diffusion with an exact solution and a sampleable, closed-form logistic-Gaussian transition kernel; Ito’s lemma gives its SDE.Source: §3.3 · p. 3
  • The logistic-normal score has a closed-form expression for score-matching training.Source: §3.3 · p. 3

3.4. Implementation Considerations

  • For k-category data, simplex-corner one-hot vectors are relaxed at the forward-process start to [α, β, …, β]; α=0.9 is suggested for k=6.Source: §3.4 · p. 3
  • At the interval boundaries, the reverse-SDE term is bounded, unlike the regular score.Source: Figure 2 · p. 3

4. Results

Initial experiments apply Simplex Diffusion to a three-category MNIST dataset and show random samples across the reverse process. Samples are converted back to discrete categories with argmax.

4. Results

  • The initial experiment maps MNIST pixel values from [0, 1, …, 255] to three categories and uses a 35-million-parameter U-Net score model.Source: §4 · p. 3
  • Generated simplex vectors are converted to discrete categories by taking their argmax.Source: §4 · p. 4
  • Figure 3 shows random samples at the beginning, middle, and end of reverse sampling with T=1000 denoising steps.Source: Figure 3 · p. 4

5. Discussion

The discussion contrasts the proposed simplex and unit-cube diffusion with prior bounded-domain methods, emphasizing their differing distributional behavior and score computation.

5.1. Simplex Diffusion

  • The proposed OU and Logistic-Normal simplex diffusion remains Logistic-Normal throughout, whereas the compared Gamma-based process has a Dirichlet limit but not during its transient dynamics.The discussion attributes the proposed process’s persistence to the correspondence between diffusion spaces in R^d and S^d via Ito’s lemma.Source: §5.1 · p. 4

5.2. Unit-Cube Diffusion

  • Reflected Diffusion targets unit-cube image generation by reflecting trajectories that would otherwise leave the domain, unlike thresholding, which disconnects training and generation.The discussion frames these methods as approaches to diffusion on bounded domains.Source: §5.2 · p. 4, §5 · p. 4
  • The proposed method extends to the unit cube by taking products of one-dimensional processes, retaining a closed-form score that is easy to implement.Reflected Diffusion’s forward score cannot be written in closed form, and its practical implementation uses two approximations.Source: §5.2 · p. 4

6. Conclusion

The conclusion presents a diffusion method for the probability simplex and unit cube with exact SDE dynamics, and identifies further evaluation as future work.

6. Conclusion

  • The proposed method performs diffusion on both the probability simplex and unit cube, with exact SDE dynamics that fit the common diffusion training paradigm.Source: §6 · p. 4
  • Future work includes testing the method on more complex datasets and evaluating its properties on the categorical distribution.Source: §6 · p. 4
  • The authors identify entropy as a possible natural notation for aleatoric uncertainty over generated values.Source: §6 · p. 4

A.1. Score Derivation

The score derivation proceeds through gradients of the second and exponential terms, then combines the resulting terms, including α and β.

A.1. Score Derivation

  • The derivation first finds the gradient of the second term, treating the log normalizing constant as having no gradient.Source: §A.1 · p. 6
  • The exponential term is handled separately before the terms are combined again.Source: §A.1 · p. 6
  • The final step combines the α and β terms to obtain the result.The preceding steps address the gradient of the second term and the exponential term before combining terms again.Source: §A.1 · pp. 6–7

A.2. Sampling and Ito’s Lemma

The section starts from an OU process and uses Ito’s Lemma to express the process in drift-and-diffusion form, including a corresponding process on the simplex.

A.2. Sampling and Ito’s Lemma

  • The OU process is split into cases j = i and j ≠ i, followed by a summary of its diffusion and drift terms.Source: §A.2 · pp. 7–8
  • The corresponding process on the simplex is introduced, with a_j given as x_j + 1/2(1 − 2x_j).Source: §A.2 · pp. 7–8
  • A self-contained definition of σ is provided for the section.Source: §A.2 · p. 7
  • Ito’s Lemma is used to write X_t in the form f(X,t)dt + G(X,t)dB_t.The passage identifies H_X as the Hessian and drops time dependence of S_t and X_t for notational simplicity.Source: §A.2 · p. 7
  • The derivation next treats the trace Hessian term.Source: §A.2 · p. 8

A.3. Diffusion Matrix Divergence

The appendix derives the divergence of the Hessian H = G(x)G(x)⊤ by expanding its coordinate-wise form and evaluating diagonal and off-diagonal cases. It then combines the resulting terms to obtain d_i.

A.3. Diffusion Matrix Divergence

  • For symmetric G(x), the Hessian is H = G(x)G(x)⊤, with a coordinate-wise form used to derive its divergence.The matrix divergence is defined similarly to the definition in Song et al. (2021).Source: §A.3 · pp. 8–9
  • The derivation analyzes G(x) in separate diagonal and off-diagonal cases before calculating the divergence.It further expands the combined expression for d_i.Source: §A.3 · pp. 9–10
  • The derivation obtains d_i by summing contributions from the rows of H and treating k = i and k ≠ i separately.It introduces helper functions a(x_i) and b(x_i) for the two cases before combining the results.Source: §A.3 · pp. 9–10

Related papers

Loading 2309.02530v2…