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

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

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.

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.

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.)

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})$

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

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

Abstract

Figures & tables

Appendix

Diffusion Models

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

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