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

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

cs.CVAug 22, 2026

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

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

Organizations: University of Cambridge, UK · King’s College London, UK · MediaTek Research, UK · University of Warwick, UK

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.

Figures & tables

Figure 1: Covariance matrices for denoising diffusion probabilistic models are well approximated by a Kronecker-DCT (K-DCT) structure. ( A ) Illustration of our K-DCT covariance approximation ( Equation 12 ), with learnable parameters indicated in green. $\otimes$ denotes the Kronecker product. ( B ) Top: covariance of the CIFAR-10 dataset (image width $d=32$ ), shown across RGB channels ( $3\times 3$ block structure) but restricted to pixels along the vertical (left) and horizontal (center) image midlines, and shown in full ( $d\times d$ block structure, right) but for the red channel only. Bottom: same visualizations of the nearest (in minimum squared error sense) approximation of the CIFAR-10 covariance that conforms to the structure shown in (A). See also Figures 5 and 6 for further analyses of how Equation 12 accurately describes the marginal covariance structure of larger datasets (ImageNet and CelebA), see Figure 11 for other data modality, speech data. ( C ) Same as (B), for an example posterior covariance matrix obtained from the Jacobian of a score network (see Section 2.1 ) pre-trained on CIFAR-10 ( 2 ) , evaluated at a partially denoised sample ( $600$ denoising steps, i.e. roughly mid-way through denoising).See also Figure 7 for a further dissection of how Equation 12 performs at various stages of denoising on CIFAR-10.

Figure 2: FID (log-scale) vs. NLL (log-scale) for different training objectives (OCM, NPR and SN) and different covariance models (heuristic, diagonal, ours) in the skip-step regime of few sampling steps ( $K$ ). Our model consistently achieves both lower FID and NLL.

Covariance	Type	Intuition
Heuristic large $\beta_{t}$ ( 15 )	Isotropic constant	Cov. of $q(\bm{x}_{t}\|\bm{x}_{t-1})$
Heuristic small $\tilde{\beta}_{t}$ ( 15 )	Isotropic constant	Cov. of $q(\bm{x}_{t}\|\bm{x}_{t-1},\bm{x}_{0})$
SN-diagonal ( 2 )	Diagonal $\bm{x}_{t}$ -dependent	Learn from data, $\mathbb{E}(\bm{\epsilon}_{t}^{2}\|\bm{x}_{t})$
NPR-diagonal ( 2 )	Diagonal $\bm{x}_{t}$ -dependent	Learn from data, Cov $(\bm{\epsilon}_{t}\|\bm{x}_{t})$
OCM-diagonal ( 29 )	Diagonal $\bm{x}_{t}$ -dependent	Learn from score, $\nabla_{\bm{x}_{t}}\log\tilde{q}(\bm{x}_{t},t)$
LowRank ( 26 )	Low-rank $\bm{x}_{t}$ -dependent	Learn from data, Cov $(\bm{\epsilon}_{t}\|\bm{x}_{t})$

Table 1: Different covariance estimation methods ranked by increasing expressiveness

FID	CIFAR10 (LS)			CIFAR10 (CS)			CelebA $64\times 64$			ImageNet $64\times 64$
# Timesteps $K$	10	25	50	10	25	50	10	25	50	25	50	100
Heuristic $\tilde{\beta}_{t}$	44.45	21.83	15.21	34.76	16.18	11.11	36.69	24.46	18.96	29.21	21.71	19.12
Heuristic $\beta_{t}$	233.41	125.05	66.28	205.31	84.71	37.35	294.79	115.69	53.39	170.28	83.86	45.04
SN-diagonal	24.06	6.91	4.63	16.33	6.05	4.17	20.60	12.00	7.88	27.58	20.74	18.04
NPR-diagonal	32.35	10.55	6.18	19.94	7.99	5.31	28.37	15.74	10.89	28.27	20.89	18.06
NPR-K-DCT (ours)	23.06	9.12	5.92	17.26	7.79	5.58	19.69	13.72	9.80	23.98	18.90	17.44

Table 2: FID score and NLL across various datasets with different sampling steps. Colors denote

{\color[rgb]{0.76,0.01,0}\textbf{1}^{\textbf{st}}}

and

{\color[rgb]{0.06,0.46,1}\textbf{2}^{\textbf{nd}}}

best (i.e. lowest) FID and NLL values.

Figure 3: Wall-clock sampling time vs. FID for different datasets and covariance models (standard heuristic, diagonal, ours). The sampling time is an average measurement for a batch size of 128 on the CIFAR10, CelebA and ImageNet datasets, and 64 on the LSUN Bedroom dataset, all on a single A6000-48GB GPU. CIFAR10, CelebA and LSUN Bedroom are tested on

K=\{10,25,50\}

steps, while ImageNet is tested on

K=\{25,50,100\}

(hence three points per curve).

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Expressiveness of the DCT covariance parameterization . Each panel

(i,j)

shows

(F\otimes F)^{\top}\text{diag}(e_{i+dj})(F\otimes F)

where

F

is the

d

-points DCT matrix, and

e_{k}

is the

k^{\text{th}}

row of the identity matrix

I_{d^{2}}

(here,

d=5

). Thus, these are the primitive basis functions from which any

(F\otimes F)^{\top}\text{diag}(\bm{\lambda})(F\otimes F)

in Equation 12 can be assembled.

Figure 5: Accuracy of our covariance parameterization for modeling the marginal (i.e. prior) ImageNet covariance . Left : normalized MSE (black) obtained when approximating the marginal ImageNet covariance $\Sigma$ by a sum of $n$ spatio-chromatic Kronecker products (“ $\text{colors}(3\times 3)\otimes\text{pixels}(d^{2}\times d^{2})$ ”). The optimal decomposition of order $n$ has an analytical form ( 42 ) . Our approximation in Equation 12 corresponds to $n=1$ , and is nearly perfect here; this best single Kronecker product approximation is shown at the top, along with the two corresponding factors. We will call the spatial factor $S\in\mathbb{R}^{d^{2}\times d^{2}}$ ). For comparison, we also show the SVD-based rank- $n$ truncation of $\Sigma$ (green). Right: normalized MSE (black) obtained when approximating the spatial factor $S$ (c.f. above) with a sum of vertical-horizontal Kronecker products ( $[d\times d]\otimes[d\times d]$ ). For $n=1$ , this type of approximation is close to, but not exactly the same, as what we propose in Equation 12 for modeling the spatial component. It would be the same if (i) we constrained our $D\equiv\text{diag}(\bm{\lambda})$ to itself be a Kronecker product of two smaller diagonals, but (ii) learned the eigenbasis instead of forcing it to be the DCT matrix $F$ . Again, a single unconstrained Kronecker product provides an excellent approximation here, indicating that a Kronecker-structured eigenbasis is empirically justified. Moreover, fixing the eigenbasis to $F$ but still learning a full diagonal $D$ (our main proposal) does nearly as well (dashed red; $(F\otimes F)^{\top}D(F\otimes F)$ ) whilst having lower computational complexity. In summary , this figure shows that the ImageNet dataset has approximately separable spatio-chromatic components , and is sufficiently translation invariant for its spatial structure to be compactly described using the DCT.

Figure 6: Same as Figure 5 , but for modelling the CelebA marginal (i.e. prior) covariance. Perhaps unsurprisingly, the spatio-chromatic structure of this particular dataset is not as separable as for ImageNet; this is likely due to certain locations in the image being dominated by certain colors (e.g. celebrities aren’t known for their blue noses). It is also less translation invariant (i.e. less diagonalizable in the DCT eigenbasis) owing to the nature of these centered portraits.

Figure 7: Accuracy of Equation 12 for modeling conditional covariances at different stages of denoising (time $t$ ) on CIFAR-10. Left : at each time, we have a “true” covariance obtained by differentiating through the score network (c.f. Section 2.1 ) shown in the first row of images. As in Figures 5 and 6 , for each of these matrices we can find its nearest “diagonal + spatio-chromatic Kronecker product” approximation (second row). This yields a fairly good reconstruction MSE (black), much better than the nearest “diagonal + rank-1” approximation (green; third row of images). Both approximations were obtained through optimization. Right : at each time, we can then focus on the spatial component of the Kronecker product identified on the left, and approximate it in the same ways as we did in Figure 5 (right) with $n=1$ . This is shown with the same color code. In summary, we see that at all times, the use of Kronecker products at both levels outperforms low-rank forms, and that reducing the complexity by restricting the spatial component to a DCT-diagonalized form does not hurt much.

Method	Step	FID $\downarrow$	AC $\uparrow$	Step	FID $\downarrow$	AC $\uparrow$
K-DCT (ours)	10	14.28	72.00 %	100	5.12	82.42%
$\Pi$ GDM ( 37 )		33.89	47.48%		4.01	84.08 %
OCM ( 29 )		77.60	36.57%		25.93	70.95%
Tweedie’s ( 34 )		-	-		75.21	55.16%

Appendix

Table 3: Inpainting+denoising results. FID and AC are reported on

50k

samples.

Panel of Figure 8 — Figure 8: Conditional posterior sampling for the noisy inpainting problem, with $\bm{y}=A\bm{x}+\sigma_{\bm{y}}\bm{\epsilon}$ , where $A$ is a random mask that covers $75\%$ of the pixels, and additional Gaussian noise of $\sigma_{\bm{y}}=10^{-3}$ . Each sub-figure shows conditional denoising sampling procedure for a different method but under the same seed, with the corrupted observations $\bm{y}$ and the original images $\bm{x}$ on the last two rows.

Datasets	Noise scheduler	Sampling steps	Pretrained model
CIFAR10	Linear	$1000$	2
CIFAR10	Cosine	$1000$	2
CelebA $64\times 64$	Linear	$1000$	38
ImageNet $64\times 64$	Cosine	$4000$	28
LSUN Bedroom $256\times 256$	Cosine	$1000$	28

Appendix

Table 4: Pretrained first-order model used in our experiments, and details of the noise schedulers

Datasets	$\text{NN}_{1}$	$\text{NN}_{2},\bm{\varepsilon}_{\phi}$	$\text{NN}_{2},C_{\phi}$	$\text{NN}_{2},\bm{d}_{\phi}$
CIFAR10 (LS)	$\text{Conv}_{\text{last}}$	$\text{Conv}_{\text{last}}$	$\text{MLP}_{\text{mid}}$	$\text{MLP}_{\text{mid}}$
CIFAR10 (CS)	$\text{Conv}_{\text{last}}$	$\text{Conv}_{\text{last}}$	$\text{MLP}_{\text{mid}}$	$\text{MLP}_{\text{mid}}$
CelebA $64\times 64$	$\text{Conv}_{\text{last}}$	$\text{Conv}_{\text{last}}$	$(\text{Res}+\text{MLP})_{\text{mid}}$	$(\text{Res}+\text{MLP})_{\text{mid}}$
ImageNet $64\times 64$	$\text{Conv}_{\text{last}}$	$(\text{Res}+\text{Conv})_{\text{last}}$	$(\text{Res}+\text{MLP})_{\text{mid}}$	$(\text{Res}+\text{MLP})_{\text{mid}}$
LSUN Bedroom $256\times 256$	$\text{Conv}_{\text{last}}$	$(\text{Res}+\text{Conv})_{\text{last}}$	$(\text{Res}+\text{MLP})_{\text{mid}}$	$(\text{Res}+\text{MLP})_{\text{mid}}$

Appendix

Table 5: Architecture details of parametric heads for the first and second-order model as in Equation 15

	Training				Sampling
	NPR		OCM
Datasets	diagonal	K-DCT	diagonal	K-DCT	diagonal	K-DCT
CIFAR10	118.74	124.82 (+ 5.1%)	310.02	318.89 (+ 2.8%)	115.74	112.49
CelebA $64\times 64$	285.25	297.37 (+ 4.2%)	680.49	727.64 (+ 6.9%)	219.71	219.47
ImageNet $64\times 64$	230.34	238.16 (+ 3.4%)	581.51	656.89 (+ 11.4%)	332.61	335.63

Appendix

Table 6: Averaged time (in millisecond) of one iteration of update w.r.t different objectives (training) and one step of denoising (sampling) for a batch size of

128

on an A6000-48GB GPU

FID	CIFAR10 (LS)			CIFAR10 (CS)
# Timesteps $K$	10	25	50	10	25	50
MEAN (K-DCT-NPR)	22.921	9.142	5.880	17.545	7.661	5.568
STD (K-DCT-NPR)	0.154	0.019	0.040	0.032	0.108	0.011
MEAN (K-DCT-OCM)	21.980	9.082	5.914	12.865	6.298	5.092
STD (K-DCT-OCM)	0.291	0.022	0.022	0.132	0.026	0.042
	CelebA $64\times 64$			ImageNet $64\times 64$

Appendix

Table 7: Mean and standard deviation of FID and NLL.

Figure 9: Scalable, efficient evaluation of the log determinant for log-likelihood computations. Left : Stochastic estimation (black; mean and quartiles from 100 independent runs) vs. exact value (red) of the log determinant involved in the log-likelihood at different stages of denoising (

t

) on CIFAR10. Here, stochastic estimation is done as shown in Equation 44 with a single probe vector

\bm{\xi}

. Center : Wall time spent in computing the log determinant (CPU). Right : Relative MSE in log-determinant estimation, estimated from 100 independent runs. Note that the terms that dominate the log-likelihood (small

t

) are also those that are approximated the best. For large

t

, the noise in Hutchinson’s trace estimator with only one probe vector is comparable to the log-det itself, hence the high relative MSE. However, these terms contribute very little to the overall NLL, and they are also the cheapest to compute, so one could use many more probe vectors if accuracy was really required there.

Covariance model	FID	Memory (MB)	Time (ms)
CIFAR10 (LS)
NPR-diag	32.35	272.39	417.66
LowRank(r=5)	28.81	320.46	419.68
LowRank(r=10)	27.57	350.52	420.52
LowRank(r=25)	26.79	440.72	425.36
LowRank(r=50)	26.14	591.05	428.15

Appendix

Table 8: Model performance and computation complexity between low-rank methods and K-DCT.

Panel of Figure 10 — Figure 10: Eigen-spectrum of the posterior covariance at various points in the sampling process. Columns from left to right correspond to time from $t=50$ to $t=T-50$ . Eigenvalues of the true posterior covariance are shown in black (obtained by evaluating the derivative of the pretrained score network at some random sample as in Section 2.1 ). Eigenvalues of the fitted covariance with (a) diagonal assumption or (b) our proposed K-DCT structure are shown in orange, where the color from yellow to orange shows the dynamics along training iterations. Within each sub-figure, the first row shows eigenvalues of covariance of the noise, Cov $[\bm{\epsilon}_{t}|\bm{x}_{t}]$ , while the second row shows eigenvalues of the skip-step covariance used during sampling, Cov $[\bm{x}_{t-k}|\bm{x}_{t}]$ , for $k=20$ . Horizontal lines in red and black indicate the value of ‘large’, $\beta_{t}$ , and ‘small’, $\tilde{\beta}_{t}$ , heuristics, respectively.

Model	FID
	$K=10$	$K=25$	$K=50$
DCT	17.51	10.88	7.57
FreeDiag	17.49	10.80	7.52
Free	16.97	10.60	7.47

Appendix

Table 9: What does the DCT matrix bring?

Figure 11: Accuracy of DCT-based parameterization for audio (speech) data. Left: Empirical covariance matrix $\Sigma$ of speech (LibriSpeech dataset) at $\tau=0.1$ ms resolution (top), $1$ ms (middle) and $5$ ms (bottom). Center: Nearest DCT-diagonalized approximation, $\Sigma\approx\hat{\Sigma}\equiv F^{\top}DF$ where $F$ is the DCT matrix and $D$ is the diagonal matrix that minimizes $\|\Sigma-\hat{\Sigma}\|_{\text{F}}^{2}$ . Right: Marginal slices along the secondary diagonal, which collapse those covariance matrices into covariance as a function of time lag $(t-t^{\prime})/\tau$ assuming stationarity. The DCT parameterization provides a very good fit to the empirical covariance at all resolutions. (Due to limited number of samples in the dataset, covariances on slower timescales ( $\tau>10$ ms) could not be estimated as accurately.)

Diffusion Models

Diffusion models learn to reverse gradual data corruption, enabling flexible generation, editing, forecasting, and inference.

See all 1,192 papers

May 13, 2026

Covariance-aware sampling for Diffusion Models

We present a covariance-aware sampler that improves the quality of pixel-space Diffusion Model (DM) sampling in the few-step regime. We hypothesize that in the few-step regime samplers fail because they rely solely on the predicted mean of the reverse distribution, while our solution explicitly models the reverse-process covariance. Our method combines Tweedie's formula to estimate the covariance with an efficient, structured Fourier-space decomposition of the covariance matrix. Implemented as an extension of DDIM, our method requires only a minimal overhead: one extra Jacobian-Vector Product (JVP) per step. We demonstrate that for pixel-based DMs, our method consistently produces superior samples compared to state-of-the-art second order samplers (Heun, DPM-Solver++) and the recent aDDIM sampler, at an identical number of function evaluations (NFE).

Andrea Schioppa, Tim Salimans

GDM - Amsterdam · work done at GDM - Amsterdam

Diffusion Model Sampling Diffusion Sampling Pixel-Space Diffusion Models Few-Step Diffusion Sampling Covariance Estimation

Ask PDF

Jun 2, 2026

Flicker-DDPM: Accelerating Denoising Diffusion via 1/f Colored Noise Injection

We propose a novel diffusion model, Flicker-DDPM, which incorporates flicker (1/f) noise inspired by self-organized criticality (SOC), a widely observed phenomenon in natural systems. Unlike denoising diffusion probabilistic models (DDPMs), which employ isotropic white noise in the forward process, Flicker-DDPM adopts colored noise with power-law spectra to better match the spectral statistics of natural images, whose power spectra typically follow P(k) proportional to 1/k^α. To this end, we develop a colored-noise module based on a spatial correlation kernel, σ(d) = (d + 1)^{-η}, and theoretically establish that adjusting η controls the spectral exponent α of the generated 1/fα noise, enabling adaptation to datasets with diverse spectral characteristics. On CIFAR-10, Flicker DDPM matches or surpasses the generation quality of a standard DDPM baseline using 3.33 times fewer sampling steps, with negligible additional computational cost per step. We further develop a frequency-domain linear theory demonstrating that spectrally matched colored noise linearizes the reverse trajectory, theoretically explaining the observed sampling acceleration.

KeXiang Mao, FanCheng Li

School of Physics and Technology, Wuhan University, Wuhan 430072, China · Hongyi Honor College, Wuhan University, Wuhan 430072, China

Few-Step Diffusion Sampling Diffusion Model Acceleration Diffusion Models

Ask PDF

May 21, 2026

The Value of Covariance Matching in Gaussian DDPMs and the Lanczos Sampler

A central error measure in Gaussian DDPMs is the path-space KL divergence between the exact reverse chain and the learned Gaussian reverse process. This quantity is especially relevant for procedures such as classifier guidance, which perturb the entire reverse trajectory rather than only the terminal sample. Prior analyses show that standard isotropic reverse covariances suffer an unavoidable

Ω(1/T)

path-KL error as the number of denoising steps

T

grows. We show that matching the full posterior covariance breaks this barrier, yielding an order-wise improvement that reduces the path KL to

O(1/T^2)

. To make full covariance matching practical, we introduce the Lanczos Gaussian sampler (LGS), a training-free, matrix-free method for sampling from the optimal reverse covariance using only covariance-vector products, which are available through Jacobian-vector products of the posterior mean. LGS avoids dense covariance storage and auxiliary covariance models. We prove that LGS approximation error decays exponentially in the number of Lanczos steps, where each Lanczos step requires a single Jacobian-vector product. Empirically, using only just three such steps improves sample quality over strong diagonal-covariance baselines, including OCM-DDPM, across standard image benchmarks. This identifies full covariance matching as both theoretically valuable and practically accessible for fast DDPM sampling.

Md Sahil Akhtar, Aymane El Gadarri, Vivek F. Farias +1

Electrical Engineering and Computer Science · Massachusetts Institute of Technology · Operations Research Center +1

Diffusion Model Sampling Distribution Matching Diffusion Models

Ask PDF

Abstract

Diffusion Models

Explore similar work

Covariance-aware sampling for Diffusion Models

Flicker-DDPM: Accelerating Denoising Diffusion via 1/f Colored Noise Injection

The Value of Covariance Matching in Gaussian DDPMs and the Lanczos Sampler