Diffusion on the Probability Simplex

TL;DR

This paper introduces a novel diffusion process on the probability simplex using the Ornstein-Uhlenbeck process and softmax, enabling better handling of discrete data and extending to bounded image generation.

Contribution

It proposes a new diffusion method on the probability simplex with a softmax-based Ornstein-Uhlenbeck process, bridging continuous and discrete data modeling.

Findings

Effective diffusion on the probability simplex for categorical data

Extension of the method to diffusion on the unit cube for image generation

Provides a natural interpretation of diffusion in terms of probability distributions

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.