Staged Diffusion Models with Analytically Designed Hyperparameters

We present StaDM, a framework for staged diffusion modeling where the generation of high-dimensional data is partitioned into multiple stages, with each stage performing diffusion on a projected subspace of dimensions. We derive an analytical objective in terms of an ideal denoiser which allows staging hyper-parameters to be optimized without expensive retraining and generation steps. For the reverse process, we design a back-projection strategy for switching between stages, thereby eliminating the need for training special bridging networks. We illustrate the usefulness of staged diffusion with (1) semi-autoregressive staging where each stage denoises a disjoint subset of dimensions chosen analytically, and (2) multi-resolution staging with analytically chosen switch points instead of existing fixed switch points. On image generation tasks we achieve up to 35\% reduction in sample generation time over homogeneous image diffusion models. The staging hyper-parameters obtained using our method are significantly faster to obtain than empirical generate-and-test approaches.