Stochastic Heat Diffusion Models

This work introduces a family of generative diffusion models based on the stochastic heat equation. The deterministic heat equation has natural connections with kernels over discrete structures and has been successfully applied in machine learning to discriminative tasks for various types of data such as graphs. We extend it to a stochastic version by formulating our diffusion as a Gaussian process whose transition kernel can be written in closed form for both the forward- and reverse-time processes. Our diffusion process propagates noise smoothly among neighbours under the guidance of a Laplacian matrix that describes the global structure of the input. Following a flow matching-based approach, we train a denoising neural network that takes noisy representations as input and gradually removes noise to generate clean samples. We experimentally show that, by exploiting the Laplacian during training, our decoder is able to recover the global structure of the input.