Improved denoising diffusion probabilistic models with efficient non-diagonal covariance modeling

AI summary

Overview

• Research area: Generative modeling with Denoising Diffusion Probabilistic Models (DDPMs), specifically the design of better denoising posterior covariance approximations to accelerate sampling while preserving likelihood-based evaluation.
• Technical level: Advanced. The paper draws on Tweedie's formula for higher-order posterior moments, Kronecker factorization, Discrete Cosine Transform (DCT) spectral theory, and Jacobian-vector products through pre-trained score networks.
• Scope in one sentence: The paper introduces a Kronecker-DCT (K-DCT) structured covariance model for image DDPMs that captures non-diagonal spatio-chromatic correlations at near-diagonal computational cost, and shows improved FID and negative log-likelihood relative to diagonal baselines in the few-step sampling regime.

What This Paper Is About

Standard DDPM sampling runs many small denoising steps because the denoising posterior covariance is only well approximated by a simple diagonal (or constant) matrix in the limit of small steps. Using fewer, larger steps requires a more accurate posterior covariance, and previous attempts modeled it either as diagonal or as diagonal-plus-low-rank — both of which ignore the strong non-diagonal correlations between neighboring pixels and between RGB channels that natural images actually exhibit. This paper's goal is to build a covariance model that respects those correlations while remaining cheap enough in memory and compute to be used at every denoising step.

Key Contributions

1. A new covariance parameterization for image DDPMs. The K-DCT model (Equation 12) combines a diagonal baseline, a full 3×3 inter-channel covariance via a Kronecker factor C_φ C_φᵀ, and an inter-pixel spatial covariance that is diagonal in the 2D DCT basis with learned positive eigenvalues λ_φ.

2. A DCT-based motivation grounded in image statistics. The authors argue that natural image covariance is not circulant (as a DFT-based model would assume) but Toeplitz-plus-Hankel, and note that matrices diagonalized by the DCT have precisely this structure. They verify empirically that the CIFAR-10 marginal covariance, in both horizontal and vertical directions, is a superposition of a Toeplitz and a Hankel component.

3. Efficient training and sampling. The Kronecker-DCT structure reduces training and sampling complexity from O(D²) to O(D log d), with an O(D) memory footprint — the size of a single image. The training objectives (OCM and NPR) and the sampling procedure are both expressed as element-wise operations plus dimension-d linear transforms, with pseudocode given in Algorithms 1, 2 and 3.

4. Empirical validation across four image datasets. Starting from pre-trained first-order UNet score models, the authors learn K-DCT-structured amortizations of the posterior covariance on CIFAR-10, Celeb-A, ImageNet (downsampled to 64×64) and LSUN Bedroom, and compare against previously reported constant, diagonal and low-rank methods.

Main Findings

• The K-DCT structure fits image covariance well: The nearest minimum-squared-error approximation conforming to the K-DCT structure provides a good fit to the marginal CIFAR-10 covariance (image width d = 32), shown across RGB channels (3×3 block structure) and in the spatial dimensions. Further analyses of the marginal covariance structure of ImageNet and CelebA are reported, as well as for speech data.

• Posterior covariances are far from diagonal when steps are skipped: The paper argues and illustrates (Figure 1C, using a score-network Jacobian evaluated at a partially denoised sample at roughly mid-way through a 600-step denoising process) that the posterior covariance is far from scalar or diagonal for realistically small horizons.

• Low-rank approximations are expensive in this setting: Capturing 90% of the variance of the CIFAR-10 denoising covariance requires a rank of r = 635, approximately 20% of 3D. The authors also cite prior work using r = 50 (6% of D) on MNIST, and note that eigenvalue spectra of image denoising posterior covariances decay slowly.

• Better FID and likelihood at few sampling steps: Across the training objectives compared (OCM, NPR and SN) and covariance models (heuristic, diagonal, and the proposed model), the K-DCT model consistently achieves both lower FID and lower negative log-likelihood in the skip-step regime of few sampling steps K.

• Negligible overhead: The additional components required by the non-diagonal model (the C_φ and d_φ heads) add little overhead relative to diagonal covariance models; an empirical comparison of training and sampling time cost between the diagonal covariance model and K-DCT is reported.

• Specific FID and NLL numbers per dataset are not contained in the provided paper content (the text is truncated within Table 2), so individual dataset-level scores cannot be reported here.

Methodology in Plain English

The authors start from an existing, already-trained diffusion model that predicts noise at each denoising step. A known mathematical relationship (an extension of Tweedie's formula) says that a well-trained noise predictor implicitly contains information about the covariance of the denoising posterior, not just its mean. The challenge is that for a color image with D = d² pixels, the full covariance matrix is 3D × 3D, which is far too large to store or use.

Their solution is to impose structure on the covariance that matches how images actually behave. Two assumptions do the work. First, correlations between color channels and correlations between spatial locations are approximately separable, so the covariance can be written as a Kronecker product of a small 3×3 color term and a large spatial term. Second, image statistics are approximately translation invariant, so the spatial term should be diagonal in a frequency basis — but the authors choose the DCT rather than the DFT, because images are bounded (not cyclic) and their covariance is Toeplitz-plus-Hankel rather than circulant.

To train the model, they add small extra prediction heads on top of the frozen pre-trained UNet — one reading from the last up-block to predict the diagonal term, and two reading from the last middle-block to predict the color factor and the DCT eigenvalues. These heads are optimized with one of two objectives: an objective that matches the covariance implied by derivatives of the score network, or a mean-squared-error objective that matches the observed noise prediction residual. Because the covariance has a tensor product structure, matrix-vector products and Frobenius norms can be computed with element-wise scaling and small (dimension-d) linear transforms rather than by ever forming the full matrix. Sampling exploits the fact that a sum of two covariance structures is a sum of two independent Gaussians, each of which has an accessible matrix square root.

Why This Matters

This work matters because it addresses a central tension in diffusion models: high sample quality traditionally requires many sequential denoising steps, but fewer, larger steps demand a posterior covariance model that diagonal and low-rank approximations cannot supply. By making a structured full-covariance approximation tractable, the paper opens a path to faster sampling that retains DDPM's tractable likelihood — an advantage that distilled one-step samplers give up.

Real-world applications:

• Faster image generation: Fewer denoising steps at comparable or better fidelity translates directly into lower inference cost for image synthesis systems.
• Image restoration: The paper notes that prior work has leveraged the link between denoising posteriors and score Jacobians to improve image restoration; a more accurate posterior covariance could improve such tasks.
• Uncertainty quantification: Because the method produces an explicit, input-dependent covariance, it provides a principled picture of denoising uncertainty rather than a heuristic constant.
• Beyond images: The authors include analysis of speech data as another modality where the same covariance structure is examined, suggesting possible extension to other signal domains.

Industry relevance: The paper's authors include researchers from MediaTek Research, and the released code builds on previous open-source work for fair comparison. Reducing the number of sequential network evaluations needed per generated image is directly relevant to deployment costs in image generation, editing and restoration products. The method's negligible per-step overhead is important, since any acceleration is negated if each step becomes much more expensive.

Future Directions

• Extending to larger and higher-resolution images. The paper notes that for larger images such as CelebA and ImageNet, better forms of approximation are needed that can capture the full rank of the posterior covariance without introducing an additional compute/memory tradeoff; how well K-DCT scales to substantially higher resolutions is not resolved in the provided content.
• Combining with latent-space diffusion and ODE-based samplers. The paper contrasts its approach with latent diffusion, deterministic ODE solvers and one-step distillation; whether the covariance structure carries over to those settings is an open question.
• Applying the model outside vision. The marginal covariance of speech data is analyzed, but the paper's main experiments are on images; extending the parameterization to other modalities is a natural next step.
• Refining the structural assumptions. The model assumes separability of chromatic and spatial correlations and approximate translation invariance; testing how robustly these hold, and what to do when they fail, remains open.

Target Audience

This paper is aimed at machine learning researchers and practitioners already familiar with diffusion models who are working on sampling acceleration, likelihood-based evaluation, or the probabilistic structure of denoising posteriors. Readers from an optimization or second-order approximation background will recognize the parallel the authors draw to curvature modeling with Kronecker-factored approximations. It is also relevant to engineers who need faster diffusion inference but do not want to give up the ability to compute likelihoods, and to researchers interested in the statistical structure of natural images.

Authors’ abstract

The sampling process of Denoising Diffusion Probabilistic Models (DDPMs) can be accelerated by leveraging second-order information in the form of approximations to the denoising posterior covariance -- allowing samples of acceptable quality to be produced in fewer but larger sampling steps. Previous attempts at using such information have used drastic (e.g.\ diagonal) simplifications of the covariance. These do not do justice to the peculiar statistical structure of natural images, which exhibit strong non-diagonal correlations between pixels and color channels, and a slow-decaying power-law frequency spectrum. Here, we develop a novel covariance model that captures these features. Our Kronecker-DCT (K-DCT) model uses a Kronecker-factored decomposition of inter-color covariances and spatial covariances modeled in the frequency domain using the Discrete Cosine Transform (DCT). The use of the DCT reduces the computational complexity from quadratic to log-linear, resulting in negligible computational and memory overhead in each denoising step. By learning K-DCT-structured amortizations of the denoising posterior covariance using pre-trained score models on CIFAR-10, Celeb-A, ImageNet and LSUN datasets, we show improved performance compared to previous SOTA denoising samplers, both in terms of FID and likelihoods, especially in the regime of few denoising steps.