Improved denoising diffusion probabilistic models with efficient non-diagonal covariance modeling
Proposes K-DCT covariance model to accelerate DDPM sampling, capturing non-diagonal correlations in natural images, improving quality with fewer steps.
Key Findings
Methodology
This work introduces a Kronecker-DCT (K-DCT) covariance structure that models inter-channel and spatial pixel correlations in natural images. By learning parameterized covariance matrices in the frequency domain, combined with pre-trained score models, the approach reduces computational complexity from quadratic to log-linear. The model captures the strong non-diagonal correlations typical in natural images, such as pixel and color channel dependencies, enabling fewer denoising steps without sacrificing quality. Extensive experiments on CIFAR-10, Celeb-A, ImageNet, and LSUN datasets demonstrate significant improvements in FID and likelihood metrics, especially in low-step regimes, with negligible overhead.
Key Results
- On CIFAR-10, the proposed K-DCT model achieves a FID of 2.8 with fewer than 10 steps, outperforming the diagonal covariance baseline at 3.3 by approximately 15%. On ImageNet, negative log-likelihood improves by 3.2 units, indicating better statistical fit. The training and sampling times increase marginally (~10%), yet the generated images exhibit sharper details and higher diversity. Ablation studies confirm that frequency domain covariance captures image statistics more accurately than low-rank or diagonal models, especially under limited steps.
- Across datasets, the model consistently outperforms baselines in quality and efficiency. The frequency domain approach effectively models long-range pixel and color correlations, leading to more realistic and detailed images. Results show that fewer steps suffice to produce high-fidelity images, making the method suitable for real-time applications. The approach's robustness is validated through extensive hyperparameter tuning and ablation experiments.
- The model's efficiency stems from FFT-based computations and tensor product structures, enabling scalable training and sampling. The learned covariance matrices provide a better inductive bias for natural images, facilitating high-quality generation with minimal computational overhead. These advances open new avenues for fast, high-resolution image synthesis and potentially extendable to video and multimodal data.
Significance
This research addresses the fundamental bottleneck in diffusion models: balancing sampling speed and image fidelity. By accurately modeling the complex, non-diagonal correlations inherent in natural images, the K-DCT approach enables high-quality, few-step generation, bridging the gap between efficiency and realism. It advances the theoretical understanding of covariance structures in generative modeling and offers practical benefits for deploying diffusion models in resource-constrained environments. The method's scalability and effectiveness could influence future research in probabilistic modeling, image synthesis, and beyond, fostering more realistic and efficient AI-generated content.
Technical Contribution
The paper introduces a novel covariance parameterization combining Kronecker decomposition and DCT in the frequency domain, enabling efficient modeling of high-dimensional, non-diagonal covariance matrices. This approach leverages the spectral properties of natural images, reducing computational complexity from O(D²) to O(D log d). It integrates seamlessly with pre-trained score networks, allowing end-to-end optimization of posterior covariance estimates. The framework supports fast training and sampling, with theoretical guarantees on approximation quality, and demonstrates superior performance in low-step diffusion sampling scenarios, setting new state-of-the-art benchmarks.
Novelty
This work is the first to integrate Kronecker factorization with DCT-based frequency domain modeling for posterior covariance in diffusion models. Unlike prior methods limited to diagonal or low-rank approximations, the proposed K-DCT captures complex pixel and channel correlations efficiently. Its ability to model long-range dependencies in natural images with minimal overhead marks a significant innovation, especially in the context of accelerating diffusion-based image synthesis with fewer steps.
Limitations
- The model relies on pre-trained score networks, which may limit generalization to unseen data distributions or highly noisy scenarios. Its assumptions about translation invariance and stationarity may not hold in highly textured or non-uniform images. Although computationally efficient, the frequency domain operations still pose challenges at ultra-high resolutions, requiring further optimization. Additionally, the approach's effectiveness in video or 3D data remains to be explored, and future work should address these extensions.
Future Work
Future research could focus on multi-scale covariance modeling to better handle non-stationary image regions. Incorporating adaptive frequency transforms or learned basis functions may enhance robustness. Extending the framework to video generation and multimodal tasks, such as text-to-image synthesis, is promising. Further optimization of FFT-based operations and hardware acceleration could enable real-time applications. Exploring unsupervised or self-supervised training paradigms to improve generalization is also a valuable direction.
AI Executive Summary
Diffusion probabilistic models have revolutionized high-fidelity image synthesis, yet their sampling efficiency remains a challenge. Traditional approaches rely on many small denoising steps, which, while effective, are computationally intensive. Recent efforts to incorporate second-order information, specifically the posterior covariance, have shown promise in reducing step count. However, existing covariance approximations are often simplistic, typically diagonal, neglecting the rich, non-diagonal correlations present in natural images.
This paper introduces a novel covariance modeling framework called Kronecker-DCT (K-DCT), which leverages the spectral properties of natural images. By decomposing the covariance into inter-channel and spatial components using Kronecker products, and modeling spatial correlations in the frequency domain via the Discrete Cosine Transform, the authors achieve a highly expressive yet computationally efficient structure. This approach captures long-range pixel and color channel dependencies, which are crucial for realistic image generation.
Through extensive experiments on datasets like CIFAR-10, Celeb-A, ImageNet, and LSUN, the authors demonstrate that their model significantly outperforms previous state-of-the-art methods, especially in regimes with fewer denoising steps. The results show improvements in FID scores, with CIFAR-10 scores dropping to 2.8 in less than 10 steps, and better likelihood estimates, confirming the model's superior statistical fidelity. The frequency domain covariance reduces computational complexity from quadratic to log-linear, enabling scalable training and sampling.
This work advances the theoretical understanding of covariance structures in diffusion models, providing a practical tool for high-quality, fast image synthesis. Its implications extend to real-time applications, large-scale content creation, and multimodal AI systems. Despite its strengths, the approach depends on pre-trained score models and assumptions about image stationarity, which may limit its applicability in certain scenarios. Future directions include multi-scale modeling, video extension, and hardware optimization, promising a broad impact across AI-generated content fields.
Deep Dive
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.