Partitioned Diffusion Priors

• Partitioned diffusion priors are probabilistic frameworks that partition high-dimensional signals into blocks, enabling scalable, flexible generative modeling with local independence.
• They leverage independent score-based diffusion processes on each partition, significantly reducing computational load in applications such as MRI, CT, and signal demixing.
• By integrating interval partition theory, Lévy processes, and Bayesian nonparametrics, these priors preserve local structure while ensuring robust inference and reconstruction.

Partitioned diffusion priors are specialized probabilistic frameworks that model complex, structured uncertainty over partitioned domains—spatial, temporal, or categorical—using generative diffusion processes defined independently or jointly on partitions. These frameworks enable both computational scalability and flexible modeling of heterogeneous, high-dimensional signals and latent structures. They have been formulated and applied across domains including scientific imaging, Bayesian nonparametrics, source separation, and function-space modeling, with mathematical constructions rooted in branching diffusions, interval partitions, score-based generative modeling, and structured transformations such as Hankel decompositions.

1. Mathematical Foundations and Generic Construction

Partitioned diffusion priors originate from the principle of decomposing a global domain (e.g., a high-dimensional signal or function, or a set of latent variables) into subsets or blocks, each equipped with an independent or weakly coupled diffusion process. In continuous domains, such as interval partitions, this entails associating a one-dimensional diffusion (with generator $$\mathcal L f(x) = b(x)f'(x) + \tfrac{1}{2}\sigma^2(x)f''(x)$$) to the width or “mass” of each interval, evolving via an Itô SDE until absorption at domain boundaries. Conditional propagation—births, splits, collapses—are governed by excursion measures (Pitman–Yor) and further ordered via spectrally positive Lévy processes, giving rise to strong Markov and branching properties in the collective state space of partitions (Buckland, 2024, Forman et al., 2019).

In high-dimensional generative modeling, the construction proceeds by splitting the full-parameter space (e.g., multi-slice volume $X \in \mathbb{C}^{S \times N \times N}$) into contiguous or arbitrary blocks $X^{(g)}$ and training independent score-based diffusion models $p_\theta^{(g)}$ or $s_\theta^{(g)}$ on each block. This partitioning can follow spatial, spectral, tensor, or feature-based decompositions, as demonstrated in medical imaging (multi-slice MRI, CT), signal demixing, and function decomposition (Valdy et al., 7 Dec 2025, Zhang et al., 2024, Wagner-Carena et al., 6 Oct 2025).

2. Algorithms, Inference Schemes, and Computational Benefits

The principal advantage of partitioning is a drastic reduction in per-block computational and memory footprint, permitting the tractable training and inference of full-scale diffusion models on commodity hardware. In multi-slice imaging applications (Valdy et al., 7 Dec 2025), the volume is split into $G$ blocks, each processed in parallel per-GPU by a backward DDPM sampler. After each timestep, blocks are aggregated for potential physics-driven updates (e.g., data consistency, proximal mapping) before the next cycle. Crucially, partitioning into contiguous blocks preserves local correlations, yielding only minor fidelity drop in generative metrics, while enabling $>\!8\times$ memory reduction.

In iterative source separation (Wagner-Carena et al., 6 Oct 2025), the joint source prior is factorized as $$q_\Theta(u) = \prod_\beta q_{\theta^\beta}(x^\beta)$$, and inference proceeds via Monte Carlo EM: joint samples are drawn from the posterior using the partitioned prior and likelihood, and each diffusion model is retrained on its assigned block.

Few-shot low-dose CT reconstruction augments partitioned priors with tensorized Hankel decompositions. The sinogram is partitioned into multiple overlapping Hankel blocks $H^s$, each modeled by its own diffusion network. Iterative stochastic solvers, equipped with low-rank projections and total variation denoising, sequentially update the partitions and fuse their outputs, yielding rapid, artifact-suppressing reconstructions even with extreme data scarcity (Zhang et al., 2024).

3. Theoretical Properties: Markovianity, Branching, and Self-Similarity

A recognized unifying aspect is the Markov (and frequently branching) nature conferred by partitioned evolution. Interval partition diffusions (Buckland, 2024, Forman et al., 2019) constitute strong Markov processes on the space of interval partitions, with block-wise independent evolution and explicit Laplace functional semigroups. The branching property ensures that future evolution is governed independently within each block, permitting hierarchical and recursive modeling (key result: Laplace semigroup propagates multiplicatively over block initializations).

Self-similarity—particularly in stable-Lévy and BESQ-based constructions—enables parameterized scaling of the process, critical for stationarity and continuous time-change operations. In functional generative models, partitioned priors enable both local adaptivity and preservation of global structure through controlled coupling at block boundaries.

4. Partitioned Diffusion Priors in Bayesian Nonparametric Models

Partitioned interval partition diffusions form nonparametric priors over evolving clusterings in domains such as [0,1], underpinning Bayesian nonparametrics via Poisson–Dirichlet laws. These models enable continuously indexed hierarchical clustering, with predictive rules codified in explicit transition kernels. Conditioning on observed data within blocks yields closed-form posterior Laplace functionals, supporting tractable, scalable posterior inference.

In more general latent variable models, partitioned diffusion priors act as flexible, learnable nonparametric priors over the components of the signal or function, facilitating effective Bayesian updates even under unknown noise or mixing conditions (Wagner-Carena et al., 6 Oct 2025).

5. Scientific Imaging Applications and Quantitative Results

Partitioned diffusion priors have demonstrated state-of-the-art performance in scientific imaging, notably for multi-slice MRI and 4D-STEM reconstruction (Valdy et al., 7 Dec 2025), and low-dose CT synthesis (Zhang et al., 2024). By replacing global volumetric priors with partitioned alternatives, these approaches yield strong improvements in in-distribution SSIM (MRI: DART 0.968 ± 0.011 vs physics-only 0.844; CT: PHD +1.2 dB PSNR over non-partitioned priors). Out-of-distribution generalization is robust: DRIFT and DART maintain high fidelity on root MRI and hexagonal crystals, confirming that partitioning sacrifices little volumetric coherence.

From a computational perspective, partitioned diffusions reduce per-GPU memory usage from 54.6 GB (vanilla 3D) to $\sim$7 GB (partitioned), and enable convergence in as few as 10 reverse-SDE steps (CT), rather than $>\!500$ for non-partitioned models.

Application	Partition Strategy	Memory Reduction	Fidelity Gains
Multi-slice MRI	Z-axis blocks (contiguous)	$\sim$8×	SSIM up to 0.968
Low-dose CT	Partitioned Hankel matrices	$\sim$10×	PSNR +1.2 dB
Source Separation	Factorized latent variables	Linear	FID down to 1.57

6. Connections to Interval Partition Diffusions and Poisson–Dirichlet Laws

Branching interval partition diffusions and Poisson–Dirichlet stationary interval partition processes (Buckland, 2024, Forman et al., 2019) provide rigorous probabilistic models for partitioned domains, with stationary (or pseudo-stationary) distributions characterized by explicit Laplace exponents and generator structure. In continuum limit, ranked block masses recover Ethier–Kurtz and Petrov’s infinite-dimensional diffusions on the Kingman simplex, and are tightly linked to pure partition priors for Bayesian nonparametrics.

These constructions yield predictive transition probabilities for block splitting and formation, support continuous hierarchical modeling, and provide explicit theoretical underpinning for partitioned prior evolution.

7. Practical Considerations, Limitations, and Extensions

While partitioned diffusion priors offer scalable and robust modeling, certain limitations and challenges persist. Key assumptions include local independence (blocks are only weakly coupled), explicit knowledge or learnability of blockwise observation operators, and appropriate data alignment. Inference at block boundaries, long-range coherence, and structure mismatches can lead to minor fidelity loss, but recent empirical results indicate high robustness in most applications. Advanced coupling strategies, overlapping block schemes, and hierarchical partitioning are active areas of extension.

A plausible implication is that partitioned diffusion priors will increasingly underpin large-scale, data-efficient generative modeling, nonparametric Bayesian clustering, and tractable high-dimensional inference, particularly in domains where global modeling is computationally infeasible or physically ill-posed.