PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rui Xia, Ayan Das, Artem Artemev, Andi Zhang, Guillaume Hennequin, Alberto Bernacchia

arXiv:2608.21972v1cs.CVstat.ML

TL;DR

DDPM samplers need efficient non-diagonal covariance models because skip-step denoising makes posterior covariance less diagonal while unrestricted models are computationally prohibitive. The paper introduces K-DCT, combining Kronecker factorization with DCT-domain spatial modeling, and learns it from pretrained score models. Across four image datasets, K-DCT improves image quality and likelihoods over diagonal models, especially with few denoising steps, although Gaussian covariance modeling may remain limited for strongly non-Gaussian skip-step posteriors.

  • Problem

    Existing posterior covariance approximations are diagonal or low-rank, despite natural images exhibiting strong non-diagonal correlations, while unrestricted covariance parameterizations require O(D^2) resources.

  • Method

    The paper learns K-DCT-structured amortizations of input-dependent denoising posterior covariance using pretrained score models.

  • Results

    K-DCT consistently improves FID and negative log-likelihood over diagonal approximations on CIFAR-10, Celeb-A, ImageNet and LSUN, especially with few sampling steps.

  • Takeaways & Limitations

    Efficiently modeling important non-diagonal natural-image statistics improves image DDPMs particularly in the few-denoising-step regime.

  • Takeaways & Limitations

    Skip-step denoising posteriors may be far from Gaussian, limiting the broader covariance-modeling approach even with an accurate covariance.

Abstract

from arXiv · show

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.

1 Introduction

DDPMs generate images through sequential denoising, creating a speed–quality trade-off when many small steps are required. This work models non-diagonal posterior covariance with K-DCT structure to improve few-step sampling efficiently.

  • Motivation: DDPMs generate images by iteratively denoising noisy states through a Markov chain, with strong fidelity and mode-coverage performance.The process starts from pure noise and repeatedly samples less noisy images.
  • Motivation: Small denoising steps make the posterior covariance approximately diagonal, but many sequential steps limit sampling parallelization and create a speed–quality trade-off.Most prior work therefore focused on estimating the posterior mean rather than its covariance.
  • Covariance structure: Existing covariance models use diagonal or low-rank structures that restrict representation of natural images’ non-diagonal spatial and chromatic correlations.The paper argues that neighboring pixels and color channels exhibit strong correlations that these simplifications miss.
  • Proposed method: K-DCT combines Kronecker-factorized color and spatial covariances represented in the Discrete Cosine Transform frequency domain.The model exploits approximate separability of chromatic and spatial correlations and approximate translation invariance.
  • Results: On CIFAR-10, Celeb-A, ImageNet and LSUN, learned K-DCT posterior covariances improve FID and negative log-likelihood over diagonal approximations, especially with few sampling steps.The amortizations are learned from pretrained score models and target input-dependent denoising posterior covariance.

2 Background

The background develops Gaussian denoising and DDPM sampling, then describes score-based and least-squares approaches for learning posterior covariance approximations. Skip-step sampling makes covariance structure especially important, while tractability constrains usable parameterizations.

  • Gaussian denoising: In Gaussian denoising, posterior moments can be obtained from score derivatives or learned as least-squares conditional estimators.The posterior mean relates to the score through Tweedie’s formula, while covariance relates to the score Jacobian.
  • Covariance learning: MMSE-based methods learn covariance amortizations by minimizing squared error, and the paper adapts this strategy alongside score derivative-based objectives for non-diagonal models.Both approaches rely on pretrained first-order noise prediction.
  • DDPMs: DDPMs construct noisy states with a forward Gaussian Markov chain and generate samples by reversing the chain from an approximately normal terminal state.Each reverse conditional is generally intractable and is commonly approximated by a Gaussian.
  • DDPMs: Skip-step DDPMs merge consecutive denoising transitions, reducing the number of sampling steps through affine relations between posterior means and covariances.Skipping steps makes the posterior covariance less diagonally dominant, increasing the importance of modeling its structure.
  • Posterior moments: Posterior means are typically estimated through a neural network that predicts effective noise and is transformed into the denoising conditional mean.The network approximates the conditional expectation of the effective noise term.
  • Posterior moments: Common posterior covariance heuristics use time-dependent scalar multiples of the identity, such as β_tI or ˜β_tI.These approximations are independent of the current noisy input x_t.
  • Covariance learning: Score derivative-based covariance learning adapts optimal covariance matching around a pretrained first-order noise predictor, while stochastic Jacobian-vector products avoid materializing non-diagonal residuals.The approach uses an unbiased stochastic gradient estimator for the covariance objective.

3 Covariance Parameterizations

The section contrasts diagonal and low-rank covariance approximations with K-DCT, which models separable color–spatial correlations efficiently. Its DCT-based structure supports tractable training and sampling with linear memory and log-linear computation.

  • Limitations of existing parameterizations: Full Cholesky parameterization requires prohibitive O(D^2) memory and compute, motivating structured covariance approximations.
  • Limitations of existing parameterizations: Diagonal covariance models have O(D) complexity but ignore pairwise pixel correlations.
  • Limitations of existing parameterizations: Capturing 90% of CIFAR-10 denoising-covariance variance requires a low-rank r = 635 ≈20% of 3D, limiting scalability to larger images.
  • K-DCT parameterization: K-DCT models approximately separable spatio-chromatic correlations through a Kronecker product of inter-channel and inter-pixel components.
  • K-DCT parameterization: The DCT is used instead of the DFT because bounded-image covariances exhibit Toeplitz-plus-Hankel structure rather than circulant structure.
  • K-DCT parameterization: The K-DCT model uses a 2D DCT spatial eigenbasis with positive parameterized eigenvalues and has O(D) memory plus O(D log d) training and sampling.
  • Efficient training: K-DCT evaluates non-diagonal covariance objectives through tensor-product matrix operations, replacing naive O(D^2) computation with O(D log d).
  • Efficient sampling: Sampling uses independent Gaussian contributions from the diagonal and Kronecker-DCT terms instead of explicitly computing the full matrix square root.

4 Experiments & Results

Experiments across multiple image datasets evaluate K-DCT covariance models against diagonal and heuristic alternatives using FID, NLL, sampling time, and structural ablations. K-DCT consistently improves likelihood and image quality, especially with fewer sampling steps, while retaining negligible overhead and outperforming low-rank alternatives.

  • Experimental setup: Experiments use pretrained first-order UNets with additional covariance heads, optimized using generalized NPR or OCM objectives.The covariance model shares most parameters with the pretrained noise predictor; two smaller components are added beyond the diagonal parameterization.
  • Experimental setup: Evaluation covers CIFAR10, CelebA, down-sampled ImageNet, and LSUN Bedroom with FID and NLL as the primary metrics.The reported comparisons include multiple noising schedules and previously reported heuristic and diagonal covariance baselines.
  • Main results: K-DCT consistently outperforms diagonal models in NLL, with larger gains when fewer sampling steps make diagonal covariance less accurate.The comparison includes both score derivative-based OCM and MMSE-based NPR training objectives.
  • Main results: FID improvements largely track the NLL gains, especially for aggressive skip-step regimes and larger images.FID values in Table 2 are computed from 50k generated samples.
  • Efficiency: K-DCT provides the best FID–NLL tradeoff among evaluated models while introducing an overhead measured against wall-clock sampling time across datasets.The timing study includes the higher-resolution 256 × 256 LSUN Bedroom dataset.
  • Ablations: K-DCT consistently outperforms diag+low-rank and purely diagonal models because posterior eigenvalues follow a slowly decaying power law that resists rank truncation.Increasing rank improves FID only up to r = 50 for CIFAR10 and r = 100 for CelebA, while memory use rises substantially for low-rank models.

5 Discussion & Limitations

K-DCT efficiently models important natural-image covariance structure and improves DDPMs particularly when few denoising steps are used. The authors note that gains remain less striking than expected and that broader covariance modeling may face Gaussian-approximation and generalizability limits.

  • K-DCT modeling of natural-image statistics improves image DDPMs especially in the regime of few denoising steps.
  • Performance gains are less striking than expected, and FID remains difficult to match against distilled models.The authors suggest skip-step denoising posteriors may be far from Gaussian, making Gaussian sampling inadequate even with the correct covariance.
  • Generalizability beyond image data remains an open question, although audio and speech with approximate translation invariance are possible application domains.The authors report a speech proof of principle and speculate that full covariance models may help even without such symmetries.

A.1 Skip-step DDPM

The skip-step DDPM formulation relates predicted noise moments to a Gaussian approximation of the image posterior q(xs|xt) for earlier time s < t.

  • Skip-step DDPM sampling approximates q(xs|xt) as a Gaussian with mean µs(xt) and covariance Σs(xt) for s < t.
  • The formulation uses standard notation for the quantities relating predicted noise moments to the skip-step posterior.

A.2 Expressiveness of the DCT parameterization

The DCT parameterization builds spatial covariance from DCT basis functions and provides accurate, efficient approximations of image covariance structure. Kronecker products outperform low-rank forms across denoising times, while DCT restriction adds little error.

  • DCT basis construction: DCT covariance bases are primitive functions assembled parametrically into the spatial component of Eq. 12.Each basis is formed by transforming a diagonal selector with F ⊗ F.
  • ImageNet covariance: A single spatio-chromatic Kronecker product nearly perfectly approximates marginal ImageNet covariance, while the DCT-based model performs nearly as well with lower complexity.
  • CelebA covariance: CelebA covariance is less separable and less translation invariant than ImageNet covariance because centered portraits associate particular locations with particular colors.
  • Conditional covariance: Across denoising times on CIFAR-10, nested Kronecker products outperform low-rank forms, and restricting the spatial component to DCT diagonalization adds little error.

A.3 Posterior sampling for inverse problems

The covariance models extend to conditional posterior sampling for linear-Gaussian inverse problems such as inpainting. K-DCT is especially useful for larger denoising steps, but exact matrix inversion can be costly and more efficient approximations remain future work.

  • Conditional posterior sampling: The method applies to linear-Gaussian inverse problems such as inpainting, using covariance-based Gaussian posterior approximations and likelihood-score guidance.
  • Role of covariance modeling: Heuristic covariance estimates can suffice for smaller steps, whereas larger steps require accurate covariance modeling, motivating K-DCT estimates of Cov[x0|xt] and the guidance Jacobian.
  • Guidance construction: The guidance term combines K-DCT covariance estimates with vector-Jacobian-product estimates of ∇xtE[x0|xt]⊤.
  • Inpainting and denoising results: In a CIFAR-10 task masking 75% of pixels with Gaussian noise σy = 10^-3, the model reports significantly better FID and classification accuracy than compared methods.FID and AC are reported on 50k samples.
  • Computational limitation: Exact matrix inversion is currently used because conjugate gradients perform poorly, making computation potentially time- and memory-consuming.A more efficient inverse approximation is left for future work.

A.4 Learning the covariance of the noise

The covariance of DDPM denoising noise can be derived from Tweedie-based relations and estimated through an MMSE objective. This provides the basis for generalized training objectives beyond diagonal covariance models.

  • Tweedie’s formula relates the covariance of the effective DDPM noise to derivatives of the denoising model.The derivation starts from the DDPM noising process and uses first- and second-order score information.
  • The noise covariance can alternatively be obtained as a least-squares MMSE estimator.This estimator avoids relying exclusively on derivatives of the noise predictor.
  • The MMSE covariance estimator yields a generalized NPR objective for training non-diagonal covariance models.The objective targets Cov(ϵt|xt) rather than only the gradient of the score.

A.5.1 Efficient training: log-linear complexity loss

The paper trains K-DCT covariance heads with full-covariance objectives while using matrix-vector products, stochastic trace estimation, and DCT operations to avoid dense covariance computation. Experiments report negligible overhead and improved sampling and likelihood performance, especially with few steps.

  • Efficient training: Full-covariance objectives are evaluated through matrix-vector products, squared Frobenius norms, and stochastic Hutchinson trace estimation.These operations avoid materializing the full residual or covariance matrix during training.
  • Efficient training: DCT-based matrix-vector products have near-linear spatial complexity, with implementations using DCT/iDCT operations and forward-mode automatic differentiation.The OCM algorithm samples Gaussian and Rademacher variables, computes covariance-model outputs, and evaluates Jacobian-vector products.
  • Model and objectives: K-DCT covariance heads output εϕ, Cϕ, and dϕ from pretrained UNet features and are trained with generalized NPR or OCM objectives.The covariance components are attached to selected middle- and up-block features while the pretrained first-order model remains fixed.
  • Efficiency: The K-DCT model adds negligible training and sampling time relative to diagonal covariance models, with substantially smaller extra memory than the original UNet.The comparison uses batch size 128 and measures one training update and one denoising step.
  • Sampling and ablations: Direct sample clipping performs better than covariance scaling for controlling K-DCT sampling noise, while DCT-basis relaxations yield only modest additional FID improvements.The clipping comparison concerns the penultimate sampling step, and the ablation reports the ordering Free > FreeDiag > original KDCT.
Loading 2608.21972v1…