Random Walks with Tweedie:
A Unified Framework for Diffusion Models

Algorithmic Template for Designing Score-based Diffusion Model Algorithms


Chicago Y. Park1, Michael T. McCann2, Cristina Garcia-Cardona3,
Brendt Wohlberg3, Ulugbek S. Kamilov1

1Computational Imaging Group, Washington University in St.Louis
2Theoretical Division, Los Alamos National Laboratory
3Computer, Computational and Statistical Sciences Division, Los Alamos National Laboratory

Paper (arxiv) Code
Banner

Figure 1: Unconditional image generation using the proposed score-based random walk framework, which decouples training and sampling to enable flexible noise schedules, step sizes, and temperature settings.


Abstract


We present a simple template for designing generative diffusion model algorithms based on an interpretation of diffusion sampling as a sequence of random walks. Score-based diffusion models are widely used to generate high-quality images. Diffusion models have also been shown to yield state-of-the-art performance in many inverse problems. While these algorithms are often surprisingly simple, the theory behind them is not, and multiple complex theoretical justifications exist in the literature. Here, we provide a simple and largely self-contained theoretical justification for score-based-diffusion models that avoids using the theory of Markov chains or reverse diffusion, instead centering the theory of random walks and Tweedie's formula. This approach leads to unified algorithmic templates for network training and sampling. In particular, these templates cleanly separate training from sampling, e.g., the noise schedule used during training need not match the one used during sampling. We show that several existing diffusion models correspond to particular choices within this template and demonstrate that other, more straightforward algorithmic choices lead to effective diffusion models. The proposed framework has the added benefit of enabling conditional sampling without any likelihood approximation.

Unified Frameworks for Diffusion Model Training and Sampling


Banner

Figure 2: Proposed templates and parameter settings for common diffusion model algorithms, including the score-based generative model (SGM), variance-exploding (VE) and variance-preserving (VP) diffusion models in SDE, and the denoising diffusion probabilistic model (DDPM). The templates unify algorithmic configurations for network training (specifying the noise scale distribution \(\sigma\), the weight function \(w(\sigma)\), and the denoising network \(r_\theta\;\)), and sampling (specifying the noise level \(\sigma_k\), step size \(\tau_k\), and temperature parameter \(\mathcal{T}_k\;\)). Here, \(\boldsymbol{s}_{\theta}\) is the pretrained score function, and \(\boldsymbol{\epsilon}_\theta\) is the pretrained neural network predicting the image's noise component.

Random Walks Framework Compatible with Any Diffusion Model


Banner

Figure 3: Uncondtional image generation with VP and VE-score-based sequence of random walks. This figure illustrates the flexibility of our sequence of random walks framework, which enables unconditional sampling with the VE score while also being compatible with the VP score under a unified framework. This implies that our framework can use any type of score without restricting the score training scheme.

Random Walks Framework for Solving Inverse Problems


Banner

Figure 3: Conditional image sampling results with the score-based sequence of random walks with sigmoid noise scheduling. We conditionally sample four images under the same setup. This figure demonstrates that our simplified framework extends beyond unconditional image synthesis to effectively solve inverse problems.

Paper


Bibtex


@article{park2024randomwalks, title={Random Walks with Tweedie: A Unified Framework for Diffusion Models}, author={Park, Chicago Y. and McCann, Michael T. and Garcia-Cardona, Cristina and Wohlberg, Brendt and Kamilov, Ulugbek S.}, journal={arXiv preprint arXiv:2411.18702}, year={2024} }