From Independent to Correlated Diffusion: Generalized Generative Modeling with Probabilistic Computers

Abstract

Diffusion models have emerged as a powerful framework for generative tasks in deep learning. They decompose generative modeling into two computational primitives: deterministic neural-network evaluation and stochastic sampling. Current implementations usually place most computation in the neural network, but diffusion as a framework allows a broader range of choices for the stochastic transition kernel. Here, we generalize the stochastic sampling component by replacing independent noise injection with Markov chain Monte Carlo (MCMC) dynamics that incorporate known interaction structure. Standard independent diffusion is recovered as a special case when couplings are set to zero. By explicitly incorporating Ising couplings into the diffusion dynamics, the noising and denoising processes exploit spatial correlations representative of the target system. The resulting framework maps naturally onto probabilistic computers (p-computers) built from probabilistic bits (p-bits), which provide orders-of-magnitude advantages in sampling throughput and energy efficiency over GPUs. We demonstrate the approach on equilibrium states of the 2D ferromagnetic Ising model and the 3D Edwards-Anderson spin glass, showing that correlated diffusion produces samples in closer agreement with MCMC reference distributions than independent diffusion. More broadly, the framework shows that p-computers can enable new classes of diffusion algorithms that exploit structured probabilistic sampling for generative modeling.

Figures (7)

• Figure 1: Hybrid p-computer/GPU inference loop for correlated diffusion. System-level view of one reverse-diffusion iteration. The current noisy configuration ($s_t$) is passed from the probabilistic computer to the neural network on the GPU, which produces per-site denoising probabilities $p=f_\theta(s_t)$. A binary clean-state estimate $\hat{s}_0$ is then sampled from these probabilities and sent back to the probabilistic computer, where it initializes an ensemble of forward Gibbs chains under the known couplings $J_{ij}$ and the scheduled inverse temperatures. Each chain produces a candidate for $s_{t-1}$, and the candidates are reweighted by the one-step likelihood $P(s_t\mid s_{t-1})$ before one is sampled to continue the reverse trajectory. Repeating this procedure produces successive reverse states $s_t \rightarrow s_{t-1} \rightarrow \cdots \rightarrow s_0$, yielding the final sample $s_0$.
• Figure 2: Probabilistic computing hardware primitives. (a) p-bits and g-bits serve as the building blocks for Bernoulli and Gaussian sampling, respectively. (b) Multi-state Potts variables with integer states in $[0, 2^N-1]$ can be realized either as native Potts units or as networks of interconnected p-bits. (c) A p-bit circuit and candidate stochastic nanodevice realizations, including sMTJs, Zener diodes, and SPADs. (d) Three tunable g-bits implemented on FPGA, each with 256 states and independently programmable means and variances.
• Figure 3: Independent and correlated diffusion noising in a p-bit framework: Illustration of the forward noising process implemented using probabilistic bits (p-bits), showing how structure-aware correlated noising generalizes the independent diffusion for a 2D ferromagnetic Ising system. (a) Independent diffusion noising: each p-bit $s_i$ is updated with $J_{ij}=0$ under a time-dependent local bias $h_i=s_i^{t-1}$, so that $I_i=s_i^{t-1}$; a decreasing inverse temperature schedule $\beta\to 0$ (Eq. \ref{['eq:pbit_update']}) progressively destroys correlations, resulting in factorized, site-wise noise injection (top row). (b) Correlated diffusion noising: each p-bit is updated from an interaction-dependent local field $I_i=\sum_j J_{ij}s_j + h_i$, where explicit couplings $J_{ij}$ introduce spatially correlated noise consistent with the underlying Ising structure (bottom row). The independent noising process in (a) is recovered as a special case of (b) by setting $J_{ij}=0$ and absorbing the noising strength into the bias term $h_i$.
• Figure 4: Reverse generative process (diffusion inference) using neural-network guidance and p-bit sampling: Schematic of one reverse step. Given the current noisy state $s_t$, a neural network computes per-site denoising probabilities $p=f_\theta(s_t)$, from which a clean-state estimate $\hat{s}_0$ is sampled. An ensemble of forward Gibbs chains initialized at $\hat{s}_0$ then produces candidates for $s_{t-1}$ on the probabilistic computer. These candidates are reweighted according to $\Pr(s_{t-1}\mid s_t,\hat{s}_0)\propto \Pr(s_t\mid s_{t-1})\,\Pr(s_{t-1}\mid \hat{s}_0)$, and one candidate is sampled to continue the reverse trajectory. Repeating this procedure yields a sequence of progressively denoised states ending at $s_0$.
• Figure 5: Reverse trajectories for independent and correlated kernels on the 2D ferromagnetic Ising model. Representative configurations (right-to-left) at diffusion timesteps $t=\{100,55,33,11,0\}$ (red/blue denote $s=+1/-1$). Top: independent (site-wise) kernel. Middle: Gibbs-sampled reference at matching timesteps. Bottom: correlated (interaction-aware) kernel. The mean magnetization $\langle m\rangle$ and mean energy per spin $\langle E\rangle/N$ are reported beneath each snapshot. The correlated kernel tracks the Gibbs reference more closely across timesteps, while the independent kernel produces spatially fragmented intermediate states and can deviate in magnetization even when energy appears comparable.