L EARNING MULTI - SCALE LOCAL CONDITIONAL PROBABILITY MODELS OF IMAGES
Anonymous authors
Paper under double-blind review
A BSTRACT
Deep neural networks can learn powerful prior probability models for images, as evidenced by the high-quality generations obtained with recent score-based diffusion methods.
But the means by which these networks capture complex global statistical structure, apparently without suffering from the curse of dimensionality, remain a mystery.
To study this, we incorporate diffusion methods into a class of local multi-scale models, which reduces dimensionality by assuming a stationary
Markov structure for wavelet coefficients conditioned on coarser-scale coefficients.
We instantiate this model using CNNs with local receptive fields, which enforce both the stationarity and Markov properties.
Global structures are captured using a CNN with receptive fields covering the entire (but small) low-pass image.
We test this model on a dataset of face images, which are highly non-stationary and contain large-scale geometric structures.
Remarkably, denoising, super-resolution, and image synthesis results all demonstrate that these structures can be captured with significantly smaller conditioning neighborhoods than required by a Markov model implemented in the pixel domain.
Our results show that score estimation for large complex images can be reduced to low-dimensional Markov conditional models across scales, which alleviates the curse of dimensionality.
I NTRODUCTION
Deep neural networks (DNNs) have produced dramatic advances in synthesizing complex images and solving inverse problems, all of which rely (at least implicitly) on prior probability models. Of particular note is the recent development of “diffusion methods” (Sohl-Dickstein et al., 2015), in which a network trained for image denoising is incorporated into an iterative algorithm to draw samples from the prior (Song & Ermon, 2019; Ho et al., 2020; Song et al., 2021), or to solve inverse problems by sampling from the posterior (Kadkhodaie & Simoncelli, 2020; Kawar et al., 2021).
The prior in these procedures is implicitly defined by the learned denoising function, which depends on the prior through the score (the gradient of the log density). But density or score estimation is notoriously difficult for high-dimensional signals because of the curse of dimensionality—worst-case data requirements grow exponentially with the data dimension.
How do neural network models manage to avoid this curse? Traditionally, density estimation is made tractable by assuming simple low-dimensional models, or structural properties that allow factorization into products of such models.
For example, the classical
Gaussian spectral model for images or sounds rests on an assumption of translation-invariance (stationarity), which guarantees factorization in the Fourier domain.
Markov random fields (Geman & Geman, 1984) assume localized conditional dependencies, which guarantees that the density can be factorized into terms acting on local, typically overlapping neighborhoods (Clifford & Hammersley, 1971).
In the context of visual images, these models have been effective in capturing local properties, but are not sufficiently powerful to capture long-range dependencies.
Multi-scale image decompositions offered a mathematical and algorithmic framework better suited for the structural properties of images (Burt & Adelson, 1983; Mallat, 2009).
The multi-scale representation facilitates handling of larger structures, and local (Markov) models have captured these probabilistically (e.g., Chambolle et al.
(1998); Buccigrossi & Simoncelli (1999); Malfait & Roose (1997);
Crouse et al.
Mihçak et al.
; Wainwright et al.
; ¸Sendur & Selesnick (2002);
Cui & Wang (2005)).
Recent work, inspired by renormalization group theory in physics, has indeed shown that probability distributions with long-range dependencies can be factorized as a product of Markov conditional probabilities over wavelet coefficients (Marchand et al., 2022).
Although the performance of these models is eclipsed by recent DNN models, the concepts on which they rest—stationarity, locality and multi-scale conditioning—are still of fundamental importance. Here, we use these tools to constrain and study a score-based diffusion model. A number of recent DNN image synthesis methods—including variational auto-encoders (Chen et al., 2018), generative adversarial networks (Gal et al., 2021) normalizing flow models (Yu et al., 2020;
Li, 2021)), and diffusion models (Ho et al., 2022; Guth et al., 2022)—use coarse-to-fine strategies, generating a sequence of images of increasing resolution, each seeded by its predecessor. With the exception of the last, these do not make explicit the underlying conditional densities, and none impose locality restrictions on their computation.
On the contrary, the stage-wise conditional sampling is typically accomplished with huge DNNs (up to billions of parameters), with global receptive fields. Here, we develop a low-dimensional probability model for images decomposed into multi-scale wavelet sub-bands.
Following the renormalization group approach, the image probability distribution is factorized as a product of conditional probabilities of its wavelet coefficients conditioned by coarser scale coefficients. We assume that these conditional probabilities are local and stationary, and hence can be captured with low-dimensional Markov models.
Each conditional score can thus be estimated with a conditional CNN (cCNN) with a small receptive field (RF).
The score of the coarse-scale low-pass band (a low-resolution version of the image) is modeled using a CNN with a global RF, enabling representation of large-scale image structures and organization.
We test this model on a dataset of face images, which present a particularly challenging example because of their long-range geometric structure.
Using a coarse-to-fine anti-diffusion strategy for drawing samples from the posterior (Kadkhodaie & Simoncelli, 2020), we evaluate the model on denoising, super-resolution, and synthesis, and show that locality and stationarity assumptions hold for conditional RF sizes as small as 9 × 9 without harming performance.
In comparison, the performance of CNNs restricted to a fixed RF size in the pixel domain dramatically degrades when the RF is made smaller than the image.
Thus, high-dimensional score estimation for images can be reduced to low-dimensional Markov conditional models, alleviating the curse of dimensionality.
M ARKOV W AVELET C ONDITIONAL M ODELS
Images are high-dimensional vectors.
Estimating an image probability distribution or its score is therefore plagued with the curse of dimensionality, unless one limits the estimation to a relatively low-dimensional model class.
This section introduces such a model class as a product of Markov conditional probabilities over multi-scale wavelet coefficients.
Markov random fields (Dobrushin, 1968; Sherrington & Kirkpatrick, 1975) define low-dimensional models by assuming that the probability distribution has local conditional dependencies over a graph, which is known a priori.
One can then factorize the probability density into a product of conditional probabilities, each defined over a small number of variables (Clifford & Hammersley, 1971).
Markov random fields have been used to model stationary texture images, with conditional dependencies within small spatial regions of the pixel image grid.
At a location u , such a Markov model assumes that the pixel value x ( u ) , conditioned by pixel values x ( v ) for v in a neighborhood of u , is independent from all pixels outside this neighborhood.
Beyond stationary textures, however, the chaining of short-range dependencies in pixel domain has proven insufficient to capture the complexity of long-range geometrical structures.
Many variants of Markov models have been proposed (Geman & Geman, 1984; Malfait & Roose, 1997; Cui & Wang, 2005), but none have been able to demonstrate performance comparable to recent deep networks while retaining a local dependency structure.
Based on the renormalization group approach in statistical physics (Wilson, 1971), new probability models are introduced in Marchand et al.
(2022), structured as a product of probabilities of wavelet coefficients conditioned on coarser-scale values, with spatially local dependencies.
These Markov conditional models have been applied to ergodic stationary physical fields, with simple conditional
Gibbs energies that are parameterized linearly.
Here, we generalize such models by parameterizing conditional Gibbs energy gradients with deep conditional convolutional neural networks having a local RF.
This yields a class of Markov wavelet conditional models that can generate complex structured images, while explicitly relying on local dependencies to reduce the model dimensionality.
[Figure]
An orthonormal wavelet transform uses a convolutional and subsampling operator W defined with conjugate mirror filters (Mallat, 2009), to iteratively compute wavelet coefficients (see Figure 1).
Let x 0 be an image of N × N pixels.
For each scale j > 1 , the operator W decomposes x j − 1 into:
[Equation]
where x j is a lower-resolution image and ¯ x j is an array of three wavelet coefficient images, each with dimensions N/ 2 j × N/ 2 j , as illustrated in Figure 1.
The inverse wavelet transform iteratively computes x j − 1 = W T (¯ x j , x j ) .
We now introduce the wavelet conditional factorization of probability models.
Since W is orthogonal, the probability density of x j − 1 is also the joint density of ( x j , ¯ x j ) .
It can be factorized by conditioning on x j :
[Equation]
This is performed J times, so that the lowest resolution image x
J is small enough, which yields:
[Equation]
The conditional distributions p (¯ x j | x j ) specify the dependencies of image details at scale j conditioned on the coarser scale values.
They may be represented with a conditional Gibbs energy:
[Equation]
where Z ( x j ) is the normalization constant for each x j .
The conditional Gibbs energies (2) have been used in the wavelet conditional renormalization group approach to obtain a stable parameterization of the probability model even at critical phase transitions, when the parameterization of the global
Gibbs energy becomes singular (Marchand et al., 2022).
Local wavelet conditional renormalization group models (Marchand et al., 2022) further impose that p (¯ x j | x j ) is a conditional Markov random field.
That is, the probability distribution of a wavelet coefficient of ¯ x j conditioned on values of x j and ¯ x j in a restricted spatial neighborhood is independent of all coefficients of ¯ x j and ¯ x outside this neighborhood (see Figure 1).
The HammersleyClifford theorem states that this Markov property is equivalent to imposing that E j can be written as a sum of potentials, which only depends upon values of ¯ x j and x j over local cliques (Clifford & Hammersley, 1971).
This decomposition substantially alleviates the curse of dimensionality, since one only needs to estimate potentials over neighborhoods of fixed size, which does not grow with the image size.
To model ergodic stationary physical fields, the local potentials of the Gibbs energy E j have been parameterized linearly by using prior physical models in Marchand et al.
We generalize Markov wavelet conditional models by parameterizing the conditional score with a cCNN having small RFs:
[Equation]
Computing the score (3) is equivalent to specifying the Gibbs model (2) without calculating the normalization constants Z ( x j ) , which are not needed in noise removal, super-resolution or image synthesis applications.
S CORE - BASED M ARKOV W AVELET C ONDITIONAL M ODELS
Score-based diffusion models have produced impressive image generation results.
To capture large-scale properties, however, these networks require RFs that encompass the entire image.
Our score-based wavelet conditional model leverages the Markov assumption to compute the score using cCNNs with small receptive fields, offering a low-dimensional parameterization of the image distribution while retaining long-range geometric structures.
Let y =
x + z be a noisy observation of clean image x ∈ R N 2 drawn from p ( x ) , with z ∼
N (0 , σ 2 Id) a sample of Gaussian white noise.
The minimum mean squared error (MMSE) estimate of the true image is well-known to be the conditional mean of the posterior:
[Equation]
This integral can be re-expressed in terms of the score:
[Equation]
This remarkable result, first introduced in Miyasawa (1961), exposes a direct and explicit relationship between the score of probability distributions and denoising (we reproduce the proof in Appendix B for completeness.)
Note that the relevant density is not the image distribution, p ( x ) , but the noisy observation density p ( y ) .
This density converges to p ( x ) as the noise variance σ 2 goes to zero.
Thanks to eq.
(5), the score can be approximated by learning a parametric mapping to estimate the denoised image, f ( y )
≈
ˆ x ( y ) .
Specifically, we can implement this mapping with a CNN, and optimize its parameters by minimizing the denoising MSE || f ( y ) − x || 2 over a large training set of images and their noise-corrupted counterparts.
Given a sufficiently expressive architecture, the denoising residual, f ( y ) − y , then provides an approximation of the variance-weighted score, σ 2 ∇ y log p
( y ) .
Also known as denoising score matching (Vincent, 2011), such denoiser-estimated score functions have been used in iterative algorithms for drawing samples from the density (Song et al., 2021;
Ho et al., 2020; Dhariwal & Nichol, 2021; Ho et al., 2022), or solving inverse problems (Kadkhodaie & Simoncelli, 2020; Cohen et al., 2021; Kawar et al., 2021; Laumont et al., 2022).
To model the conditional wavelet distribution p (¯ x j | x j ) , Guth et al.
(2022) parameterize the score ∇ ¯ y j log p (¯ y j | x j ) of noisy wavelet coefficients ¯ y j conditionally to a clean low-pass image
x j with a cCNN (eq.
Specifically, the cCNN takes three noisy wavelet detail channels as input, along with the corresponding low-pass channel, and generates three estimated detail channels.
The network is trained to minimize mean square error between ¯ x j and f j (¯ y j , x j ) .
Thanks to a conditional extension of eq.
(5), an optimal network computes f j (¯ y j , x j ) = ∇ ¯ y j log p (¯ y j | x j ) . Additionally, at the coarsest scale J , a CNN denoiser f J ( y J ) is trained to estimate the score of the low-pass band, ∇ y J log p ( x J ) by minimizing mean square error between x J and f J ( y J ) .
The following theorem proves that the Markov wavelet conditional property is equivalent to imposing that the cCNN RFs are restricted to the conditioning neighborhoods.
The RF of a given element of the network response is defined as the set of input image pixels on which this element depends.
Theorem 1.
The wavelet conditional density p (¯ x j | x j ) is Markovian over a family of conditioning neighborhoods if and only if the conditional score ∇ ¯ x j log p (¯ x j | x j ) can be computed with a network whose RFs are included in these conditioning neighborhoods.
The proof of the theorem is provided in Appendix A. Note that even if the conditional distribution of clean wavelet coefficients p (¯ x | x j ) satisfies a local Markov property, the noisy distribution p (¯ y j | x j ) is in general not Markovian.
However, we shall parameterize the scores with a cCNN with small RFs and hence show that both the noisy and clean distributions are Markovian.
At each scale 1 ≤ j ≤ J , the cCNN has RF that is localized in both ¯ y j and x j , and has a fixed size over all scales, independent of the original image size.
From Theorem 1, this defines the Markov conditioning neighborhoods of the learned model.
The effect of the RF size is examined in the numerical experiments of Section 4.
[Figure] Parameterizing the score with a convolutional network further implies that the conditional probability p (¯ x j | x j ) is stationary on the wavelet grid at the scale j .
However, the RF of the low-pass CNN is designed to cover the entire image x J , and thus does not enforce local Markov conditioning nor stationarity.
The product density (1) is therefore not stationary.
Despite these strong simplifications, we shall see that these stationary Markov wavelet conditional models are able to capture complex long-range image dependencies in highly non-stationary image ensembles such as centered faces.
M ARKOV W AVELET C ONDITIONAL D ENOISING
We now evaluate our Markov wavelet conditional model on a denoising task.
We use the previously learned CNNs to define a multi-scale denoising architecture, illustrated in Figure 2.
The wavelet transform of the input noisy image y is computed up to a coarse-scale J .
The coarsest scale image is denoised by applying the denoising CNN learned previously: ˆ x J = f J ( y J ) .
Then for J ≥ j ≥ 1 , we compute the denoised wavelet coefficients conditioned on the previously estimated coarse image: ˆ¯ x j = f (¯ y j , ˆ x j ) .
We then recover a denoised image at the next finer scale by applying an inverse wavelet transform: ˆ x j − 1 = W T (ˆ¯ x j , ˆ x j ) .
At the finest scale we recover a denoised image ˆ x .
Because of the orthogonality of W , the global MSE can be decomposed into a sum of wavelet MSE at each scale, plus the coarsest scale error:
∥ x − ˆ x ∥ 2 = (cid:80)
J − 1 j =1 ∥
¯ x j − ˆ¯ x j ∥ 2 + ∥ x J − ˆ x J ∥ 2 .
The global
MSE thus summarizes the precision of the score models computed over all scales.
We evaluate the peak signal-to-noise ratio (PSNR) of the denoised image as a function of the noise level, expressed as the PSNR of the noisy image.
We use the CelebA dataset (Liu et al., 2015) at 160 × 160 resolution.
We use the simplest of all orthogonal wavelet decompositions, the Haar wavelet, constructed from separable filters that compute averages and differences of adjacent pairs of pixel values (Haar, 1910).
All denoisers are “blind” (they are not informed of the noise level) and “universal” (they can operate on images contaminated with noise of any standard deviation σ ∈ [0 , 1] ).
They all have the same depth and layer widths, and their receptive field is controlled by changing the convolutional kernel size of each layer.
Appendix C provides architecture and training details.
Figure 3 shows that the multi-scale denoiser based on a conditional wavelet Markov model outperforms a conventional denoiser that implements a Markov probability model in the pixel domain.
More precisely, we observe that when the Markov structure is defined over image pixels, the performance degrades considerably with smaller RFs (Figure 3, left panel), especially at high noise levels (low PSNR).
Images thus contain long-range global dependencies that cannot be captured by Markov models that are localized within the pixel grid.
On the other hand, multi-scale denoising performance remains nearly the same for RF sizes down to 9 × 9 , and degrades for 5 × 5 RFs only at low noise levels (high PSNR) (Figure 3, right panel).
This is remarkable considering that, for the finest scale, the 9 × 9 RF implies conditioning on one percent of the coefficients.
The wavelet conditional score model thus successfully captures long-range image dependencies, even with small Markov neighborhoods.
38 33 28 23 18 13 7 2 - 3psnr in        p s n r o u t
RF = 43x43 RF = 23x23 RF = 13x
{+RF = 9x9 RF = 5x
+}
Identity 38 33 28 23 18 13 7 2 - 3psnr in
RF = 43x43 RF = 23x23 RF = 13x
RF = 9x9 RF = 5x5 only low pass Identity It is also worth noting that in the high noise regime (i.e., low PSNR), all multi-scale denoisers (even with RF as small as 5 × 5 ) significantly outperforms the conventional denoiser with the largest tested RF size ( 43 × 43 ).
The dependency on RF in this regime demonstrates the inadequacy of local modeling in the pixel domain.
On the contrary, the effective neighborhoods of the multi-scale denoiser are spatially global, but operate with spatially-varying resolution.
Specifically, neighborhoods are of fixed size at each scale, but due to the subsampling, cover larger proportions of the image at coarser scales.
The CNN applied to the coarsest low-pass band (scale J ) is spatially global, and the denoising of this band alone explains the performance at the highest noise levels (magenta curve, Figure 3).
To further illustrate this point, consider the denoising examples shown in Figure 4.
Since all denoisers are bias-free, they are piecewise linear (as opposed to piecewise affine), providing some interpretability (Mohan* et al., 2020).
Specifically, each denoised pixel is computed as an adaptively weighted sum over the noise input pixels.
The last panels show the equivalent adaptive linear filter that was used to estimate the pixel marked by the green square, which can be estimated from the associated row of the Jacobian.
The top row shows denoising results of a conventional CNN denoiser for images that are the size of the network RF.
Despite very heavy noise levels, the denoiser exploits the global structure of the image, and produces a result approximating the clean image.
The second row shows the results after training the same denoiser architecture on much larger images.
Now the adaptive filter is much smaller than the image, and the denoiser solution fails to capture the global structure of the face.
Finally, the last row shows that the multi-scale wavelet conditional denoiser can successfully approximate the global structure of a face despite the extreme levels of noise.
Removing high levels of noise requires knowledge of global structure.
In our multi-scale conditional denoiser, this is achieved by setting the RF size of the low-pass denoiser equal to the entire low-pass image size, similarly to the denoiser shown on the top row.
Then, each successive conditioning stage provides information at a finer resolution, over ever smaller RFs relative to the coefficient grid.
The adaptive filter shown in the last column is foveated : the estimate of the marked pixel depends on all pixels in the noisy image, but those that are farther away are only included within averages over larger blocks.
Thus, imposing locality in wavelet domain lifts the curse of dimensionality without loss of performance, as opposed to a locality (Markov) assumption in the pixel domain.
M ARKOV W AVELET C ONDITIONAL S UPER - RESOLUTION AND S YNTHESIS
We generate samples from the learned wavelet conditional distributions in order to visually assess the quality of the model in performing a super-resolution task.
We compare this with solving the super-resolution inverse problem directly using a CNN denoiser operating in the pixel domain.
We also compare the models on image synthesis.
We first give a high-level description of our conditional generation algorithm.
The low-resolution image x J is used to conditionally generate wavelet coefficients ¯ x J from the conditional distribution p (¯ x J | x J ) .
An inverse wavelet transform next recovers a higher-resolution image x
J − 1 from both x J and ¯ x J .
The conditional generation and wavelet reconstruction are repeated J times, increasing the resolution of the sample at each step.
In the end, we obtain a sample x from the full-resolution image distribution conditioned on the starting low-resolution image p ( x | x J ) .
x is thus a stochastic super-resolution estimate of x J .
To draw samples from the distributions p (¯ x j | x j ) implicitly embedded in the wavelet conditional denoisers, we use the algorithm defined in Kadkhodaie & Simoncelli (2021), which performs gradient ascent on the log-probability obtained from the cCNN using eq.
This is similar to score-based diffusion algorithms (Song & Ermon, 2019; Song et al., 2021;
Ho et al., 2020), but the timestep hyper-parameters require essentially no tuning, since stepsizes are automatically obtained from the magnitude of the estimated score.
Extension to the conditional case is straightforward.
The sampling algorithm is detailed in Appendix D. All the cCCN denoisers have RF 13 × 13 .
Train and test images are 320 × 320 images from the CelebA HQ dataset (Karras et al., 2018).
Because the wavelet conditional denoiser has a small RF size, the samples are generated from a model with Markov wavelet conditional distributions.
We compare these with samples from a model with a local Markov assumption in the pixel domain.
This is done using a CNN with a 40 × 40 RF trained to denoise full-resolution images, which approximates the score ∇ log p ( x ) .
Given the same low-pass image x J , we can generate samples from p ( x | x J ) by viewing this as sampling from the image distribution of x constrained by a linear measurements x J .
This is done with the same sampling algorithm of Kadkhodaie & Simoncelli (2021), with a small modification, again described in Appendix D.
For visual evaluation, Figure 5 shows super-resolution samples from these two learned image models.
The local Markov model in the pixel domain generates details that are sharp but artifactual and incoherent over long spatial distances.
On the other hand, the Markov wavelet conditional model produces much more natural-looking face images.
This demonstrates the validity of our model: although these face images are not stationary (they have global structures shared across the dataset), and are not Markov in the pixel domain (there are clearly long-range dependencies that operate across the image), the details can be captured with local stationary Markov wavelet conditional distributions.
We also tested behavior of the Markov wavelet conditional model on image synthesis.
We synthesize a 40 × 40 terminal low-pass image using the score, ∇ x J log p ( x J ) , obtained from the low-pass CNN denoiser with global RF.
Unlike the conditional wavelet stages, this architectural choice does not enforce any local Markov structure nor stationarity.
The RF is also larger than the scale of global non-stationary structures, such as the overall face shape.
The synthesis then proceeds using the same coarse-to-fine steps as used for super-resolution: wavelet coefficients at each successive scale are generated by drawing a sample using the cCNN conditioned on the previous scale.
The first (smallest) image in Figure 6 is generated from the low-pass denoiser (see Appendix D for the sampling algorithm).
We can see that it successfully captures the coarse structure of a face, made possible because the denoiser is non-stationary, with a global RF covering the extent of the low-pass image.
This image is then refined by application of successive stages of the multi-scale super-resolution synthesis algorithm described above.
The next three images in Figure 6 show successively higher resolution images generated in the procedure.
For comparison, the last image in Figure 6 shows a sample generated using a conventional CNN denoiser.
Once again, this illustrates that assuming spatially localized Markov property on the pixel grid and estimating the score with a CNN with RF smaller than the image fails to capture the non-stationary distribution of faces.
Specifically, the model is unable to generate structures larger than the RF size, and the samples are texture-like and composed of local regions resembling face parts.
We note that the quality of the generated images is not on par with the most recent score-based diffusion methods (Ramesh et al., 2022; Saharia et al., 2022; Rombach et al., 2022).
These methods have also been used for iterative super-resolution of an initial coarse-scale sample.
The sample is generated with a very large network (more than a billion parameters, much larger than ours, with 600k parameters in the low-pass CNN and 200k in the cCNNs), and each scale-wise denoiser is itself a U-Net, with associated RF that is the entire image.
Thus, the implicit probability models is global, and it is an open question whether and how these architectures are able to escape the curse of dimensionality.
Our local conditional Markov assumptions provide a step towards the goal of making explicit the probability model and its factorization into low-dimensional components.
D ISCUSSION
We have generalized a Markov wavelet conditional probability model of image distributions, and developed an explicit implementation using cCNNs to estimate the conditional model scores.
The resulting conditional wavelet distributions are stationary and Markov over neighborhoods corresponding to the cCNN receptive fields.
The coarse-scale low-pass band is modeled using the score estimated with a CNN with global receptive fields.
We trained this model on a dataset of face images, which are non-stationary with large-scale geometric features.
We find that the model, even with relatively small cCNN RFs, succeeds in capturing these features, producing high-quality results on denoising and super-resolution tasks.
We contrast this with local Markov models in the pixel domain, which are not able to capture these features, and are instead limited to stationary ergodic textures.
The Markov wavelet conditional model demonstrates that probability distributions of images can be factorised as products of conditional distributions that are local.
This model provides a mathematical justification which can partly explain the success of coarse-to-fine diffusion synthesis (Ho et al., 2020; Guth et al., 2022), which also compute conditional scores at each scale.
Although we set out to understand how factorization of density models could allow them to avoid the curse of dimensionality in training, it is worth noting that the dimensionality of the conditioning neighborhoods in our network is still uncomfortably high ( 4 × 9 × 9 = 324 ).
Although this is reduced by a factor of roughly 300 relative to the image size ( 320 × 320 = 102 , 400 ), and this dimensionality remains constant even if this image size is increased, it is still not sufficient to explain how the conditional score can be trained with realistic amounts of data.
In addition, the terminal low-pass CNN operates globally (dimensionality 40 × 40 = 1600 ).
Thus, the question of how to further reduce the overall dimensionality of the model remains open.
Our experiments were performed on cropped and centered face images, which present a particular challenge given their obvious non-stationarity.
The conditional models are approximately stationary due to the fully convolutional structure of the cCNN operations (although this is partially violated by zero-padded boundary handling).
As such, the non-stationary capabilities of the full model arise primarily from the terminal low-pass CNN, which uses spatially global RFs.
We speculate that for more diverse image datasets (e.g., a large set of natural images), a much larger capacity low-pass CNN will be needed to capture global structures.
This is consistent with current deep networks that generate high-quality synthetic images using extremely large networks (Ramesh et al., 2022; Saharia et al., 2022; Rombach et al., 2022).
On the other hand, the cCNNs in our model all share the same architecture and local RFs, and may (empirically) be capturing similar local conditional structure at each scale.
Forcing these conditional densities to be the same at each scale (through weight sharing of the corresponding cCNNs) would impose a scale invariance assumption on the overall model.
This would further reduce model complexity, and enable synthesis and inference on images of size well beyond that of the training set.
A P ROOF OF T HEOREM
To simplify notation we drop the j subscript.
Let I (resp.
J ) denote the set of indices of pixel values of ¯ x (resp.
x ).
If S is a set of indices we denote ¯ x ( S ) =
(¯ x ( i ))
i ∈ S ∩ I .
Let G be a graph whose nodes are I ∪ J .
For each i ∈ I , let N ( i ) ⊆
I ∪ J be the neighborhood of node i , with i
̸∈ N ( i ) and
N + ( i ) =
N ( i ) ∪ { i } .
To prove Theorem we need to show that the local Markov property:
[Equation]
6) is equivalent to the conditional score being computable with RFs restricted to neighborhoods:
[Equation]
for some functions f i .
We first prove that eq.
(6) implies eq.
Let i ∈ I .
We have the following factorization of the probability distribution:
[Equation]
where we have used eq.
(6) in the last step.
Then, taking the logarithm and differentiating, only the first term remains:
[Equation]
which proves eq.
Reciprocally [Equation] we now prove that eq.
(7) implies eq.
Let i ∈ I [Equation] and δ i ( j ) =
δ ij .
We have, by integrating the partial derivative:
[Equation]
[Equation]
where we have used eq.
(7) in the last step.
Note that the first term does not depend on ¯ x ( i ) , while the second term only depends on ¯ x ( N + ( i )) and x ( N + ( i )) .
This implies that when we condition on ¯ x ( N ( i )) and x ( N + ( i )) , the density factorizes as a term which only involves ¯ x ( i ) and a term which does not involve ¯ x ( i ) .
This further implies conditional independence and thus eq.
B P ROOF OF EQUATION
Miyasawa’s result, while surprising, is easily proven (Miyasawa, 1961).
The observation distribution, p ( y ) is achieved by marginalizing p ( y , x ) :
[Equation]
where the noise distribution g ( z ) is assumed Gaussian.
The gradient of the observation density of
Eq.
(B) is then:
[Equation]
Multiplying both sides by σ 2 /p ( y ) and separating the right side into two terms gives:
[Equation]
C D ETAILS OF TRAINING AND ARCHITECTURE
Architecture .
The terminal low-pass CNN and all cCNNs are bias-free.
A convolutional neural net can be made bias-free by removing all additive constants from convolution operations and batchnormalizations (i.e., division by standard deviation without subtracting the mean)
(Mohan* et al., 2020).
All networks contain 21 convolutional layers with no subsampling, each consisting ofchannels.
Each layer, except for first and last, is followed by a ReLU non-linearity and bias-free batch-normalization.
Thus, the transformation is both homogeneous (of order 1) and translationinvariant (apart from handling of boundaries), at each scale.
All convolutional kernels in the low-pass CNN are 3 × 3 , resulting in a 43 × 43 RF and 665 , 856 parameters in total.
Convolutional kernels in the cCNNs are adjusted to achieve different RF sizes.
For example, a 13 × 13 RF arises from choosing 3 × 3 kernels in every 4 th layer and 1 × 1 (i.e., pointwise linear combinations across all channels) for the rest, resulting in a total of 214 , 144 parameters.
For comparison, we also trained conventional (non-multi-scale) CNNs for denoising.
For RF 43 × 43 , we used the same architecture as for the coarsest scale band of the multi-scale denoiser: 21 bias-free convolutional layers with no subsampling.
To create smaller RFs, we followed the same strategy of setting some filter sizes in the intermesiate layer to 1 × 1 .
Training . For experiments shown in Figure 3 and Figure 4, we use 202 , 499 training and 100 test images of resolution 160 × 160 from the CelebA dataset (Liu et al., 2015).
For experiments shown in
Figure 5, Figure 8 and Figure 6, we use 29 , 900 train and 100 test images, drawn from the CelebA HQ dataset (Karras et al., 2018) at 320 × 320 resolution.
We follow the training procedure described in (Mohan* et al., 2020).
The network is trained to denoise images corrupted by i.i.d.
Gaussian noise with standard deviations drawn from the range [0, 1] (relative to image intensity range [0, 1]).
Training is carried out on batches of size 512. Note that all denoisers all universal and blind. That is, they are trained to handle a range of noise, and the noise level is not an input to the denoiser, respectively. D W AVELET C ONDITIONAL S YNTHESIS A LGORITHM
Sampling from both the CNN and cCNN denoisers is achieved using a slightly modified version of the algorithm of Kadkhodaie & Simoncelli (2021), as defined in Algorithm 1.
This method uses only three hyperparameters, and their settings are more forgiving than those of backward SDE discretization parameters in score-based diffusions.
They consist of initial and final noise levels, as well as a step size parameter h ∈ [0 , 1] which controls the trade-off between computational efficiency and visual quality.
For the examples in Figure 5, Figure 8 and Figure 6, we chose h = 0 .
01 , σ 0 = 1 , and σ ∞ is set to 0 . 0002 .
Image synthesis is initialized with a terminal low-pass image (either sampled from the associated
CNN, or computed from a test image), and successively sampling from the wavelet conditional distributions at each scale, as defined in Algorithm 2.
Require: denoiser f , step size h , initial noise level σ 0 , final noise level σ ∞ 1: t =2: Draw x 0 ∼ N (0 , σ 20 Id) 3: while σ t ≥ σ ∞ do 4: t ← t +5: d t ← f ( x t − 1 ) − x t −▷ Compute the score from the denoiser residual 6: σ 2 t ← || d t || 2 /N
▷ Compute the current noise level for stopping criterion 7: x t ← x t − 1 + h d t ▷ Perform a partial denoiser step to remove a fraction of the noise 8: end while 9: return x t
Require: number of scales J , low-pass image x J , conditional denoisers ( f j ) 1 ≤ j ≤ J , step size h , initial noise level σ 0 , final noise level σ ∞ 1: for j ∈ { J, . .
, 1 } do 2: ¯ x j ← DrawSample( f j ( · , x j ) , h, σ 0 , σ ∞ )
▷ Wavelet conditional sampling 3: x j − 1 ← W T (¯ x j , x j ) ▷ Wavelet reconstruction 4: end for 5: return x
E A DDITIONAL S UPER - RESOLUTION E XAMPLES