Discrete Spatial Diffusion: Intensity-Preserving Diffusion Modeling

TL;DR

Discrete Spatial Diffusion (DSD) reframes diffusion modeling on a discrete 2D lattice as a continuous-time, discrete-state Markov jump process that exactly conserves total intensity per color channel, addressing the limitation of traditional diffusion models in handling constrained discrete quantities. The forward process redistributes intensity units via a spatial random walk with 4-neighbor hops, while a neural network learns the time-dependent reverse-rate rules $\bar{r}$ to reconstruct samples, using either rate-matching or likelihood losses and adaptive CFL-guided tau-leaping for efficient sampling. DSD demonstrates competitive image synthesis and inpainting on standard benchmarks under exact intensity conditioning and extends to scientifically meaningful domains, generating realistic porous rock microstructures and Li-ion battery electrode morphologies with preserved porosity and phase fractions, validated by domain-specific metrics such as two-point correlations, pore-size distributions, and interface/TPB lengths. The approach enables physics-informed data generation with exact mass/volume constraints and suggests a new class of conservation-aware diffusion models with broad implications for scientific workflows and constrained generative modeling.

Abstract

Generative diffusion models have achieved remarkable success in producing high-quality images. However, these models typically operate in continuous intensity spaces, diffusing independently across pixels and color channels. As a result, they are fundamentally ill-suited for applications involving inherently discrete quantities-such as particle counts or material units-that are constrained by strict conservation laws like mass conservation, limiting their applicability in scientific workflows. To address this limitation, we propose Discrete Spatial Diffusion (DSD), a framework based on a continuous-time, discrete-state jump stochastic process that operates directly in discrete spatial domains while strictly preserving particle counts in both forward and reverse diffusion processes. By using spatial diffusion to achieve particle conservation, we introduce stochasticity naturally through a discrete formulation. We demonstrate the expressive flexibility of DSD by performing image synthesis, class conditioning, and image inpainting across standard image benchmarks, while exactly conditioning total image intensity. We validate DSD on two challenging scientific applications: porous rock microstructures and lithium-ion battery electrodes, demonstrating its ability to generate structurally realistic samples under strict mass conservation constraints, with quantitative evaluation using state-of-the-art metrics for transport and electrochemical performance.

Figures (19)

• Figure 1: Schematic diagrams illustrating how intensity is modeled in different diffusion frameworks. (a) Gaussian Diffusion relies on the Ornstein--Uhlenbeck process in continuous intensity space. (b) Prior discrete-state diffusion models apply a discrete-state Markov process to independent pixel intensities. (c) Discrete Spatial Diffusion (this work) relies on a Markov jump process of the intensity units over a discrete spatial lattice, redistributing particles while exactly conserving total intensity per color channel.
• Figure 2: (a) The forward processes for Gaussian Diffusion hoDenoisingDiffusionProbabilistic2020, Inverse Heat Dissipation Model heat_eq, and Discrete Spatial Diffusion (ours) applied on an image, sampled at discrete times. (b) Percentage change in intensity relative to the original image under the forward process.
• Figure 3: (a) Inpainting realizations on MNIST; 15% difference of conditioning intensity between consecutive rows. (b) Conditioned MNIST generations across different intensities and classes. (c) Unconditional CIFAR-10 generations. (d) Unconditional CelebA generations.
• Figure 4: Distributions of porosity and generated samples for three rock classes. Each plot shows the porosity distribution across the dataset (shaded curve) and corresponding DSD-generated microstructures at selected quantiles. (Left) Berea Sandstone, (center) Savonnières Carbonate, and (right) Massangis Limestone. Samples illustrate the model’s ability to generate structurally diverse images conditioned on exact global intensity. Quantitative metrics are presented in Appendix \ref{['app:geology_metrics']}.
• Figure 5: Generated cathode microstructures with varying phase volume fractions. The carbon binder domain appears in black, active material particles in gray, and electrolyte-filled pore space in white.