Diffusion on the Probability Simplex

TL;DR

Diffusion on the Probability Simplex introduces a diffusion framework that operates on the probability simplex to model discrete data with continuous-noise dynamics. It combines the logistic-normal distribution with an Ornstein-Uhlenbeck forward process, yielding a closed-form logistic-Gaussian transition kernel and a derived SDE on the simplex. The method extends naturally to the unit cube and preserves a closed-form score, enabling stable, likelihood-based training. Preliminary MNIST-style experiments illustrate the approach and the discussion situates it relative to prior simplex and unit-cube diffusion work.

Abstract

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.

Figures (3)

• Figure 1: Examples of the Logistic-Normal distribution (PDF values) on $\mathcal{S}^3$ with parameters $\mu = [0,0],~[0.2,0.35]$ and $\sigma = [0.5,0,5],~[0.6,0.8]$ respectively.
• Figure 2: A comparison between the regular score, $\nabla_x\textrm{log }p_t(x)$, and the reverse SDE term, $g^2(x,t)\nabla_x\textrm{log }p_t(x)$, in the one-dimensional case. The reverse SDE term is bounded at the border of the interval $[0,1]$, unlike the score. The PDF of the logistic-normal distribution is plotted for clarity, along with a dotted line around the score for visual clarity.
• Figure 3: Random samples from a Simplex Diffusion model. Samples are taken at the beginning, middle and end of the reverse process and correspond to the top middle and bottom row respectively. Sampling is done with $T=1000$ denoising steps