📊 statistics

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

This paper introduces the Kronecker-DCT (K-DCT) model, a novel non-diagonal covariance approximation that leverages Kronecker factorization and the Discrete Cosine Transform to efficiently capture natural image correlations, thereby significantly improving the quality and likelihood of Denoising Diffusion Probabilistic Models (DDPMs) with fewer sampling steps while maintaining negligible computational overhead.

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

Published 2026-08-25

📖 6 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a machine that learns to create art not by painting from scratch, but by starting with a canvas covered in static and slowly, carefully removing the noise until a clear picture emerges. This is the core idea behind a family of artificial intelligence models known as denoising diffusion models. They have become a standard tool for generating high-quality images, from realistic faces to landscapes, because they produce results that are both sharp and diverse. However, there is a catch: the process of removing the noise is inherently slow. To get a good picture, the computer must take thousands of tiny, cautious steps, peeling away a little bit of static at a time. If it tries to take larger steps to speed things up, the image often becomes blurry or distorted because the machine loses track of how the different parts of the picture relate to one another.

For a long time, researchers have tried to solve this speed problem by making a simplifying assumption: they assumed that the uncertainty in the image at any given moment could be described by looking at each pixel independently. It is a bit like trying to describe the weather by only looking at the temperature in one specific spot, ignoring how the wind in that spot affects the rain in the next. While this approach works well enough when taking tiny steps, it breaks down when researchers try to take the larger, faster steps needed for practical use. The reality of natural images is much more complex; the color of one pixel is deeply connected to its neighbors, and the red, green, and blue channels of an image influence each other in intricate ways. A new study published in the Transactions on Machine Learning Research suggests that by acknowledging these complex connections, we can make these image generators significantly faster without sacrificing quality.

The researchers, a team from institutions including the University of Cambridge and King's College London, set out to build a better map of these connections. They focused on the "posterior covariance," a technical term that essentially describes how the uncertainty of one part of the image depends on the uncertainty of another. In the old, simpler models, this map was drawn as a grid where only the diagonal lines mattered, implying that a pixel's future was independent of its neighbors. The new team argued that this was a drastic oversimplification that ignored the true structure of natural images, which exhibit strong, non-diagonal correlations between pixels and color channels. They also noted that these images have a specific frequency pattern, where details fade out in a predictable way, similar to how a musical chord has a specific structure of notes.

To capture this complexity without slowing the computer down, the team developed a new mathematical model they call K-DCT. The name comes from two key ideas they combined. First, they recognized that the relationship between colors (like red, green, and blue) and the relationship between spatial locations (left, right, up, down) are roughly separate from each other. This allowed them to break the massive, complex problem into two smaller, manageable pieces. Second, they realized that the spatial relationships in an image are best understood not by looking at the pixels directly, but by looking at the image as a collection of waves, using a technique called the Discrete Cosine Transform. This is a method that breaks an image down into its fundamental frequencies, much like how a prism breaks white light into a rainbow of colors. By working in this frequency domain, the researchers could describe the complex web of connections between pixels using a very compact set of numbers.

The result is a model that is both highly expressive and incredibly efficient. While a full description of the connections in a high-resolution image would require a massive amount of memory and computing power, the K-DCT model manages to do the same job with a footprint that is barely larger than the image itself. The researchers tested this approach on several well-known image datasets, including CIFAR-10, which contains small images of everyday objects, and CelebA, which features portraits of celebrities. They also tested it on larger, more complex datasets like ImageNet and LSUN. In every case, they compared their new method against the best existing techniques, which relied on the simpler, diagonal assumptions.

The findings were clear and consistent. When the researchers forced the models to generate images in fewer steps—simulating a scenario where speed is critical—the new K-DCT model produced significantly better results. The images were sharper, more detailed, and statistically closer to the real data they were trained on. In technical terms, the models achieved lower error rates and better likelihood scores, meaning the computer was more confident in its creations. Perhaps most importantly, the speed of the generation process did not suffer. The extra calculations required by the new model were so efficient that they added almost no time to the process, even on high-resolution images. This suggests that the bottleneck in generating images quickly was not the computer's raw power, but the way the uncertainty was being modeled.

The team also explored why this approach works so well. They found that the connections between pixels in a denoising image do not fade away quickly; instead, they follow a slow, predictable decay that spans many orders of magnitude. Simple models that try to ignore these long-range connections or approximate them with a few basic patterns fail to capture this nuance. The K-DCT model, by contrast, naturally accommodates this structure. Interestingly, the model performed exceptionally well even on datasets like CelebA, where the images are of human faces and do not have the perfect symmetry or translation invariance that the mathematical theory assumes. This suggests that the model is robust enough to handle real-world imperfections, capturing the essential "full" structure of the data even when the underlying patterns are not perfectly regular.

While the study does not claim to have solved every problem in artificial intelligence, it offers a compelling solution to a specific and persistent hurdle: the trade-off between speed and quality in image generation. By moving away from the assumption that pixels act alone and embracing the reality that they are deeply interconnected, the researchers have shown that we can accelerate the creative process of these machines. The work implies that the future of efficient image generation lies not in building bigger computers, but in building smarter models that understand the true, complex nature of the images they are trying to create. The K-DCT model stands as a proof that a more accurate understanding of statistical structure can lead to tangible improvements in performance, making high-quality image generation faster and more accessible without requiring a leap in hardware capabilities.

Technical Summary: Improved Denoising Diffusion Probabilistic Models with Efficient Non-Diagonal Covariance Modeling

Problem Statement

Denoising Diffusion Probabilistic Models (DDPMs) achieve state-of-the-art performance in image generation but suffer from slow sampling speeds due to their inherently sequential nature. While the sampling process can be accelerated by taking fewer, larger steps (skip-step sampling), this approach degrades sample quality if the denoising posterior covariance is not accurately modeled.

Previous attempts to accelerate sampling by leveraging second-order information have relied on drastic simplifications, such as assuming the posterior covariance is diagonal or low-rank. The authors argue these assumptions are insufficient because natural images exhibit strong non-diagonal correlations between pixels and color channels, as well as a slow-decaying power-law frequency spectrum. Consequently, diagonal approximations fail to capture the complex statistical structure of the denoising posterior, particularly when the number of sampling steps is small. Furthermore, modeling the full covariance matrix is computationally prohibitive ( $O(D^2)$ complexity for $D$ pixels), creating a trade-off between accuracy and tractability similar to that found in second-order neural network optimization.

Methodology

The paper proposes a novel covariance model called Kronecker-DCT (K-DCT) that captures the specific statistical properties of natural images while maintaining computational efficiency.

1. The K-DCT Covariance Model

The proposed model, $E_\phi(x_t, t)$ , approximates the denoising posterior covariance (or the covariance of the noise prediction error) using a structured decomposition:
$E_\phi(x_t, t) = \text{diag}(\varepsilon_\phi) + C_\phi C_\phi^\top \otimes (F^\top \otimes F^\top) \text{diag}(\lambda_\phi) (F \otimes F)$

The model consists of three key components:

Diagonal Baseline: A term $\text{diag}(\varepsilon_\phi)$ that absorbs any diagonal contributions, ensuring the model is at least as expressive as previous diagonal models.
Inter-Channel Correlation: Captured by the term $C_\phi C_\phi^\top \in \mathbb{R}^{3\times3}$ , which models the full correlation between RGB color channels.
Inter-Pixel (Spatial) Correlation: Modeled using a Kronecker product of the channel correlation and a spatial component. The spatial component is diagonalized in the frequency domain using the Discrete Cosine Transform (DCT) matrix $F$ $F$ .
- The authors justify the use of DCT over the Discrete Fourier Transform (DFT) by noting that natural images have Toeplitz + Hankel covariance structures (due to boundary conditions) rather than the circulant structure assumed by DFT. The DCT basis naturally diagonalizes this structure.
- The eigenvalues $\lambda_\phi$ are learned parameters in the frequency domain.

2. Computational Efficiency

The K-DCT structure allows for efficient training and sampling without explicitly constructing the full $3D \times 3D$ covariance matrix:

Complexity: The computational complexity is reduced from $O(D^2)$ to $O(D \log d)$ (where $d$ is the image dimension), making the overhead negligible compared to the forward pass of the score network.
Training: The authors adapt two existing objectives for learning covariance parameters:
- Optimal Covariance Matching (OCM): Based on the Jacobian of the score function (Tweedie's 2nd-order formula).
- Noise Prediction Residual (NPR): Based on a Minimum Mean Squared Error (MMSE) formulation.
- Both objectives are extended to handle non-diagonal models using stochastic trace estimators and efficient matrix-vector products derived from the Kronecker-DCT structure.
Sampling: Sampling is performed by adding two independent Gaussian samples: one from the diagonal term and one from the Kronecker-DCT term. The latter is sampled efficiently using the DCT basis and learned eigenvalues.

3. Training Strategy

The K-DCT model is trained as an "amortization" on top of pre-trained first-order score networks (UNets). The architecture shares parameters with the pre-trained model, with additional small neural network heads ( $NN_2$ ) dedicated to predicting the covariance components ( $\varepsilon_\phi, C_\phi, \lambda_\phi$ ). This approach requires training only a small fraction of the total parameters.

Key Contributions

Novel Covariance Parameterization: Introduction of the K-DCT model, which combines Kronecker-factored inter-channel correlations with DCT-based spatial correlations. This captures the slow-decaying power-law spectrum and non-diagonal structure of natural images more accurately than diagonal or low-rank approximations.
Efficient Implementation: Demonstration that the K-DCT model can be trained and sampled with $O(D \log d)$ complexity, avoiding the memory and compute bottlenecks of full covariance estimation.
Empirical Validation: Systematic evaluation across multiple datasets (CIFAR-10, CelebA, ImageNet, LSUN) showing that K-DCT outperforms previous state-of-the-art (SOTA) samplers, particularly in the regime of few sampling steps.

Results

The authors evaluated their method against heuristic, diagonal, and low-rank covariance models using two primary metrics: Fréchet Inception Distance (FID) for sample quality and Negative Log-Likelihood (NLL) for statistical fit.

Performance: The K-DCT model consistently achieved lower FID and lower NLL compared to diagonal and low-rank baselines across all datasets and training objectives (OCM and NPR).
Few-Step Regime: The improvements were most significant when the number of sampling steps ( $K$ ) was small (e.g., 10–25 steps). In these regimes, diagonal approximations fail to capture the necessary posterior structure, whereas K-DCT maintains high fidelity.
Trade-off: The method achieves a superior trade-off between FID and NLL compared to existing methods. While distilled one-step models may offer faster generation, they lack tractable likelihood estimation, a key advantage retained by the K-DCT approach.
Overhead: Experiments confirmed that the K-DCT parameterization introduces negligible wall-clock time overhead compared to diagonal models, even on high-resolution datasets like LSUN Bedroom (256x256).

Significance and Claims

The paper claims that modeling the non-diagonal structure of denoising posteriors is crucial for accelerating diffusion sampling without sacrificing quality. The significance of the work lies in:

Bridging the Gap: It bridges the gap between the theoretical need for accurate second-order information and the practical constraints of memory and computation.
Statistical Structure: It highlights that natural images possess specific statistical properties (separable spatio-chromatic correlations and power-law spectra) that can be efficiently exploited via the K-DCT structure.
Limitations and Future Work: The authors modestly note that while K-DCT improves performance, it does not yet match the FID of distilled one-step models. They suggest this may be because skip-step posteriors are far from Gaussian, implying that even a perfect covariance model cannot fully correct the sampling error if the Gaussian assumption itself is violated. They also speculate that while the method is designed for images, it could apply to other domains with approximate translation invariance, such as audio and speech, citing preliminary proof-of-concept results on speech data.