M ULTIMEASUREMENT G ENERATIVE M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
We formally map the problem of sampling from an unknown distribution with density p X in R d to the problem of learning and sampling p Y in R Md obtained by convolving p X with a ﬁxed factorial kernel: p Y is referred to as M-density and the factorial kernel as multimeasurement noise model (MNM).
The M-density is smoother than p X , easier to learn and sample from, yet for large M the two problems are mathematically equivalent since X can be estimated exactly given
Y =
y using the Bayes estimator (cid:98) x ( y ) =
E [ X | Y =
y ] .
To formulate the problem, we derive (cid:98) x
( y ) for Poisson and Gaussian MNMs expressed in closed form in terms of unnormalized p Y .
This leads to a simple least-squares objective for learning parametric energy and score functions.
We present various parametrization schemes of interest, including one in which studying Gaussian M-densities directly leads to multidenoising autoencoders —this is the ﬁrst theoretical connection made between denoising autoencoders and empirical Bayes in the literature.
Samples from p X are obtained by walk-jump sampling (Saremi & Hyv ¨ arinen, 2019) via underdamped
Langevin MCMC (walk) to sample from p Y and the multimeasurement Bayes estimation of X (jump).
We study permutation invariant Gaussian M-densities on
MNIST, CIFAR-10, and FFHQ-256 datasets, and demonstrate the effectiveness of this framework for realizing fast-mixing stable Markov chains in high dimensions.
I NTRODUCTION
Consider a collection of i.i.d.
samples { x i } ni =1 , assumed to have been drawn from an unknown distribution with density p X in R d .
An important problem in probabilistic modeling is the task of drawing independent samples from p X , which has numerous potential applications.
This problem is typically approached in two phases: approximating p X , and drawing samples from the approximated density.
In unnormalized models the ﬁrst phase is approached by learning the energy function f X associated with the Gibbs distribution p X ∝ exp( − f X ) , and for the second phase one must resort to
Markov chain Monte Carlo methods, such as Langevin MCMC, which are typically very slow to mix in high dimensions.
MCMC sampling is considered an “art” and we do not have black box samplers that converge fast and are stable for complex (natural) distributions.
The source of the problem is mainly attributed to the fact that the energy functions of interest are typically highly nonconvex.
A broad sketch of our solution to this problem is to model a smoother density in an M-fold expanded space.
The new density p ( y ) , called M-density , is deﬁned in R Md , where the bold y is a shorthand for ( y 1 , . . .
, y M ) .
M-density is smoother in the sense that its marginals p m ( y m ) are obtained by the convolution p m ( y m )
= (cid:82) p m ( y m | x ) p ( x ) dx with a smoothing kernel p m ( y m | x ) which for most of the paper we take to be the isotropic Gaussian:
[Equation]
Although we bypass learning p ( x ) , the new formalism allows for generating samples from p ( x ) since
X can be estimated exactly given Y =
y (for large M ).
To give a physical picture, the approach here is based on “taking apart” the complex manifold where the random variable X is concentrated in R d and mapping it to a smoother manifold in R Md where Y =
( Y 1 , . . .
, Y M ) is concentrated.
Smoothing a density with a kernel is a technique in nonparametric density estimation that goes back to Parzen (1962).
In kernel density estimation, the estimator of p ( x ) is obtained by convolving the empirical measure with a kernel.
In that methodology, the kernel bandwidth ( σ , for Gaussian kernels) is adjusted to estimate p ( x ) in R d given a collection of independent samples { x i } ni =1 .
This estimator,
But what if the kernel bandwidth is ﬁxed : how much easier is the problem of estimating p ( y ) ?
This question is answered in (Goldfeld et al., 2020), where they obtained the rate of convergence e O ( d ) n − 1 / 2 (measured using various distances) in remarkable contrast to the well-known n − 1 /d rate for estimating p ( x ) .
This nonparametric estimation result is not directly relevant here, but it formalizes the intuition that learning p ( y )
= (cid:82) p ( y | x ) p ( x ) dx is a lot simpler than learning p ( x ) .
With this motivation, we start with an introduction to the problem of learning unnormalized p ( y ) , based on independent samples from p ( x ) .
This problem was formulated by Vincent (2011) using score matching (Hyv¨arinen, 2005).
It was approached recently with the more fundamental methodology of empirical Bayes (Saremi & Hyv¨arinen, 2019).
The idea is to use the Bayes estimator of X given
Y =
y , the study of which is at the root of the empirical Bayes approach to statistics (Robbins, 1956), in a least-squares objective.
This machinery builds on the fact that the estimator (cid:98) x ( y ) =
E [ X | Y =
y ] can be expressed in closed form in terms of unnormalized p ( y ) (Sec. 3.1).
For Gaussian kernels, (cid:98) x ( y ) is expressed in terms of the score function ∇ log p ( y ) (Miyasawa, 1961).
For such kernels, the learning objective arrived at in (Saremi & Hyv¨arinen, 2019) is identical to the denoising score matching formulation (Vincent, 2011), but with new insights rooted in empirical Bayes which is the statistical framework for denoising.
The main problem with the empirical Bayes methodology is that p ( x | y ) remains unknown and cannot be sampled from.
The estimator (cid:98) x ( y ) =
E [ X | Y =
y ] can be computed, but the concentration of the posterior p ( x | y ) around the mean is not in our control.
Our solution to this problem starts with an observation that is very intuitive from a Bayesian perspective: one can sharpen the posterior by simply taking more independent noisy measurements.
This scheme is formalized by replacing p ( y | x ) with the factorial kernel
[Equation]
which we name multimeasurement noise model (MNM).
Now, the object of interest is a different density which we call M-density obtained by convolving p ( x ) with the factorial kernel:
[Equation]
This formally maps the original problem of drawing samples from p ( x ) to drawing samples from p ( y ) for any ﬁxed noise level since the estimator of X given Y =
y is asymptotically exact.
We quantify this for Gaussian MNMs using the plug-in estimator (the empirical mean of measurements).
Smooth & Symmetric!
Consider Gaussian MNMs with equal noise level σ in the regime of large σ , large M such that σ (cid:112) d/M is “small”.
In that regime, the complex manifold associated with the data distribution is mapped to a very smooth symmetric manifold in a much higher dimensional space.
The original manifold can be reconstructed via a single step by computing (cid:98) x
( y ) .
Due to equal noise levels, the manifold associated with M-density is symmetric under the permutation group:
[Equation]
where π is a permutation of indices (Fig. 1).
Although we develop a general methodology for studying
M-densities, in the later part of the paper we focus on permutation invariant Gaussian M-densities.
The paper is organized as follows.
In Sec. 2, we derive Bayes estimators for Poisson and Gaussian
MNMs.
In Sec. 3, we present the least-squares objective for learning Gaussian M-densities.
We also give a weaker formulation of the learning objective based on score matching.
Sec. 4 is devoted to the important topic of parametrization, where we introduce multidenoising autoencoder (MDAE) in which we formally connect M-densities to the DAE literature.
DAEs have never been studied for factorial kernels and the emergence of MDAE as a generative model should be of wide interest.
In addition, we introduce metaencoder formulated in an unnormalized latent variable model, which is mainly left as a side contribution.
In Sec. 5, we present the sampling algorithm used in the paper.
In Sec. 6, we present our experiments on MNIST, CIFAR-10, and FFHQ-256 datasets which were focused on permutation invariant M-densities.
The experiments are mainly of qualitative nature demonstrating the effectiveness of this method in generating fast mixing Markov chains in high dimensions.
Related works are discussed in Sec. 7, and we ﬁnish with concluding remarks.
[Figure]
Notation.
The subscripts are dropped from densities and energy functions when it is clear from their arguments: p ( y ) =
p Y ( y ) , p ( y | x ) =
p Y | X = x ( y ) , f ( y ) =
f Y ( y ) , etc.
Bold fonts are reserved for multimeasurement random variables:
Y =
( Y 1 , . . .
, Y M ) .
The following are shorthand notations: [ M ] = { 1 , . . .
, M } and ∇ m = ∇ y m .
Throughout, ∇ is the gradient with respect to inputs (in R Md ), not parameters.
The following convention is used regarding parametric functions: f θ ( · ) =
f ( · ; θ ) .
Different parametrization schemes come with a different set of parameters, the collection of which we denote by θ .
For all the datasets used in the paper, X takes values in the hypercube [0 , 1] d .
F ORMALISM : M ULTIMEASUREMENT B AYES E STIMATORS
This work is based on generalizing the empirical Bayes methodology to MNMs.
It is well known that the least-squares estimator of X given Y =
y (for any noise model) is the Bayes estimator:
[Equation]
Next we study this estimator ` a la Robbins (1956) for Poisson (the Poisson kernel was the ﬁrst example studied in 1956) and Gaussian MNMs.
In both cases the estimator (cid:98) x ( y ) is derived to be a functional of the joint density p ( y ) .
In addition, (cid:98) x
( y ) is invariant to scaling p ( y ) by a constant, therefore one can ignore the partition function in this estimation problem .
This is the main appeal of this formalism.
Poisson MNMs are included as a warm-up and to demonstrate the generality of the new formalism, but we will not pursue it as a generative model in our experiments for technical reasons due to the challenges regarding sampling discrete distributions in high dimenions (see Remark 7 ).
2. P OISSON MNM
Let X be a random variable taking values in R + .
The Poisson MNM is deﬁned by:
[Equation]
The numerator in r.h.s.
of Eq.
4 is computed next.
The measurement index m below is an arbitrary index in [ M ] used for absorbing x such that xp ( y | x ) has the same functional form as p ( y | x ) :
[Equation]
where 1 m is deﬁned as a vector whose component l is δ ml .
Using Eq.
4, it immediately follows
[Equation]
We emphasize that the dependency of (cid:98) x ( y ) on the noise channel (measurement index) m that appears on the right hand side of the expression above is only an artifact of the calculation (we observe this again for Gaussian MNMs).
The result above holds for any measurement index m ∈
[ M ] , therefore ( y m + 1) p
( y + 1 m )
=
( y m (cid:48) + 1)
p
( y + 1 m (cid:48) )
for all m, m (cid:48) ∈
[ M ] .
Example.
We can derive the estimator (cid:98) x
( y ) analytically for p ( x ) =
e − x .
We ﬁrst derive p ( y ) :
[Equation]
where the sums/products are over the measurement indices
l ∈ [ M ] .
Using Eq.
5, it follows
[Equation]
As expected one arrives at the same result by computing Eq.
5 for any measurement index m .
2. G AUSSIAN MNM
Let X be a random variable in R d .
The Gaussian MNM is deﬁned by:
[Equation]
( y
[Equation]
| x ) , where p m ( y m | x ) =
N ( y m ; x, σ 2 m I d ) .
[Equation]
It follows (as in the Poisson MNM, m in the equation below is an arbitrary index in [ M ] ):
[Equation]
We multiply both sides of the above expression by p ( x ) and integrate over x .
The derivative ∇ m (with respect to y m ) and the integration over x commute, and using Eq. 4 it follows
[Equation]
which we simplify by dividing both sides by p ( y ) :
[Equation]
This expression is the generalization of the known result due to Miyasawa (1961).
As in the Poisson
MNM, the result above holds for any m ∈ [ M ] , therefore:
[Equation]
Example.
We also studied the M-density for Gaussian MNMs analytically.
The calculations are insightful and give more intuitions on the new formalism.
We refer to Appendix A for the results.
Remark 1 ( (cid:98) x ( m ) θ ( y ) notation) .
For parametric M-densities, Eq. 8 is in general an approximation.
We use the superscript m to emphasize that there are M ways to compute the estimator:
[Equation]
Remark 2 ( σ ⊗ M notation) .
Gaussian MNMs with the constant noise level σ are denoted by σ ⊗ M .
For σ ⊗ M models, the M-density p ( y ) (score function ∇ log p ( y ) ) is invariant (equivariant) with respect to the permutation group π :
[ M ] → [ M ] .
See Fig.
1 for an illustration.
2. C ONCENTRATION OF THE PLUG - IN ESTIMATOR
The plug-in estimator of X is the empirical mean of the multiple noisy measurements we denote by (cid:98)
x mean ( y ) =
M − 1 (cid:80) m y m .
This estimator is highly suboptimal but its analysis in high dimensions for Gaussian MNMs is useful.
Due to the concentration of measure phenomenon (Tao, 2012) we have
[Equation]
where σ eff (“eff” is for effective) is given by
[Equation]
The calculation is straightforward since y m =
x + ε m where ε m ∼ N (0 , σ 2 m I d ) .
It follows: x − (cid:98) x mean ( y ) =
− M − 1 (cid:80) m ε m which has the same law as N (0 , σ 2 eff I d ) .
This calculation shows that the estimator of X in R d concentrates at the true value at a worst-case rate O ( (cid:112) d/M ) (consider replacing the sum in Eq.
10 by Mσ 2max ).
This analysis is very conservative, as it ignores the correlations between the components of y in R Md (within and across noise channels), and one expects a (much) tighter concentration for the optimal estimator.
In the next section, we present an algorithm to learn the multimeasurement Bayes estimator based on independent samples from p ( x ) .
L EARNING G AUSSIAN M- DENSITIES 3. N EURAL E MPIRICAL B AYES
In this section, we focus on learning Gaussian M-densities using the empirical Bayes formalism.
We closely follow the approach taken by Saremi & Hyv¨arinen (2019) and extend it to M-densities.
The power of empirical Bayes lies in the fact that it is formulated in the absence of any clean data.
This is reﬂected by the fact that (cid:98) x ( y ) is expressed in terms of p ( y ) which can in principle be estimated without observing samples from p ( x ) (Robbins, 1956); we generalized that to factorial kernels in
Sec. 2.
What if we start with independent samples { x i } ni =1 from p ( x ) and our goal is to learn p ( y ) ?
A key insight in neural empirical Bayes was that the empirical Bayes machinery can be turned on its head (in a form of a Gedankenexperiment presented in that paper).
We can draw samples (indexed by j ) from the factorial kernel y ij ∼ p
( y | x i ) and feed the noisy data to the empirical Bayes “experimenter” (the word used in 1956).
The experimenter’s task (our task!) is to estimate X , but since we observe X = x
i , the squared (cid:96) 2 norm (as a reminder, the Bayes estimator is the least-squares estimator) (cid:13)(cid:13)
x
i − (cid:98) x θ ( y ij ) (cid:13)(cid:13) 2 2 serves as a signal to learn p ( y ) , and also (cid:98) x
( y ) for unseen noisy data.
Next, we present the least-squares objective to learn p ( y )
∝ e − f ( y ) for Gaussian M-densities.
It is important to note that (cid:98) x
( y ) is expressed in terms of unnormalized p ( y ) , without which we must estimate the partition function (or its gradient) during learning.
Here, we can choose any expressive family of functions to parametrize (cid:98) x ( y ) which is key to the success of this framework.
Using our formalism (Sec. 2.2), the Bayes estimator takes the following parametric form (see Remark 1):
[Equation]
There are therefore M least-squares learning objectives
[Equation]
that in principle need to be minimized simultaneously, since as a corollary of Eq.
8 we have:
[Equation]
The balance between the M losses can be enforced during learning by using a softmax-weighted sum of them in the learning objective, effectively weighing the higher losses more in each update.
However, in our parametrization schemes, coming up, simply taking the sum of the M losses as the learning objective proved to be sufﬁcient for a balanced learning across all the noise channels:
[Equation]
The above learning objective is the one we use in the remainder of the paper.
3. D ENOISING S CORE M ATCHING
One can also study M-densities using score matching with the following objective (Hyv ¨ arinen, 2005):
[Equation]
In Appendix B , we show that the score matching learning objective is equal (up to an additive constant independent of θ ) to the following multimeasurement denoising score matching (MDSM) objective:
[Equation]
This is a simple extension of the result by Vincent (2011) to Gaussian MNMs.
The MDSM objective and the empirical Bayes’ (Eq. 14) are identical (up to a multiplicative constant) for σ ⊗ M models.
Remark 3.
Compared to neural empirical Bayes (NEB), denoising score matching (DSM) takes a very different approach regarding learning M-densities.
DSM starts with score matching (Eq. 15).
NEB starts with deriving the Bayes estimator of X given Y =
y for a known kernel p ( y | x ) (Eq. 7).
NEB is a stronger formulation in the sense that two goals are achieved at once: learning M-density and learning (cid:98) x ( y ) .
What remains unknown in DSM is the latter, and knowing the estimator is key here.
Without it, we cannot draw a formal equivalence between p X and p Y (see Sec. 1).
We return to this discussion at the end of the paper from a different angle with regards to denoising autoencoders.
P ARAMETRIZATION S CHEMES
We present three parametrization schemes for modeling Gaussian M-densities.
Due to our interest in σ ⊗ M models we switch to a lighter notation.
Here, the learning objective (Eq. 14) takes the form
[Equation]
) where x ⊗ M denotes ( x, . .
. , x ) repeated M times.
Parametrization schemes fall under two general groups: multimeasurement energy model (MEM) and multimeasurement score model (MSM).
In what follows, MDAE is an instance of MSM, MEM 2 & MUVB are (closely related) instances of MEM.
4. MDAE ( MULTIDENOISING AUTOENCODER )
The rigorous approach to learning Gaussian M-densities is to parametrize the energy function f .
In that parametrization, ∇ f θ is computed with automatic differentiation and used in the objective (Eq. 17).
That is a computational burden, but it comes with a major advantage as the learned score function is guaranteed to be a gradient ﬁeld .
The direct parametrization of ∇ f is problematic due to this requirement analyzed by (Saremi, 2019); see also (Salimans & Ho, 2021) for a recent discussion.
Putting that debate aside, in MDAE we parametrize the score function explicitly g θ :
R Md → R Md .
Then, g θ replaces −∇ f θ in Eq.
In particular, we consider the following reparametrization:
[Equation]
(18) simply motivated by the fact we would like to cancel y in Eq.
17, otherwise the loss starts at very high values at the initialization.
This is especially so in the regime of large σ , M , and d .
It follows:
[Equation]
This is very intriguing and it is worth taking a moment to examine this result:
modeling M-densities is now formally mapped to denoising multimeasurement noisy data explicitly .
It is a DAE loss with a multimeasurement twist.
Crucially, due to our empirical Bayes formulation, the MDAE output ν θ ( y ) becomes a parametrization of the Bayes estimator(s) (combine Eq.
18 and Eq.
[Equation] [Equation]
This makes a strong theoretical connection between our generalization of empirical Bayes to factorial kernels and a new learning paradigm under multidenoising autoencoders , valid for any noise level σ .
4. MEM 2 ( MEM with a QUADRATIC form with an optional METAENCODER )
Can we write down an energy function associated with the score function in Eq.
The answer is no, since g θ is not a gradient ﬁeld in general, but we can try the following ( θ = ( η, ζ ) ):
[Equation]
Ignoring h ζ for now, by calculating −∇ f θ we do get both terms in Eq. 18, plus other terms.
The function h ζ which we call metaencoder is optional here.
Intuitively, it captures the “higher order interactions” between y and ν , beyond the quadratic term (see below for another motivation).
The metaencoder is implicitly parametrized by η (via ν η ), while having its own set of parameters ζ .
4. MUVB ( MULTIMEASUREMENT UNNORMALIZED VARIATIONAL BAYES )
The expression above (Eq. 21) is a simpliﬁcation of an unnormalized latent variable model that one can set up, where we take the variational free energy to be the energy function.
This builds on recent studies towards bringing together empirical Bayes and variational Bayes (Saremi, 2020a;b).
We outline the big picture here and refer to Appendix C for details.
Latent variable models operate on the principle of maximum likelihood, where one is obliged to have normalized models.
Since our model is unnormalized we consider setting up a latent variable model with unnormalized conditional density.
Essentially both terms in Eq.
21 arise by considering the following ( z are the latent variables):
[Equation]
named metalikelihood which further underscores the fact that it is unnormalized.
The full expression for the energy function also involves the posterior q φ ( z | y ) .
As a remark, note that what is shared between all three parametrization schemes is ν , although vastly different in how (cid:98) x
( y ) is computed.
S AMPLING A LGORITHM
Our sampling algorithm is an adaptation of walk-jump sampling (WJS)
(Saremi & Hyv¨arinen, 2019).
We run an MCMC algorithm to sample M-density by generating a Markov chain of multimeasurement noisy samples.
This is the walk phase of WJS schematized by the dashed arrows in Fig.
At arbitrary discrete time k , clean samples are generated by simply computing (cid:98) x
( y k ) ( θ is dropped for a clean notation).
This is the jump phase of WJS schematized by the solid arrow in Fig.
What is appealing about multimeasurement generative models is the fact that for large M , p ( x | y k ) is highly concentrated around its mean (cid:98) x ( y k ) , therefore this scheme is an exact sampling scheme—this was in fact our original motivation to study M-densities.
For sampling the M-density (the walk phase), we consider
Langevin MCMC algorithms that are based on discretizing the underdamped Langevin diffusion:
[Equation]
Here γ is the friction, u the inverse mass, and B t the standard Brownian motion in R Md .
Discretizing the Langevin diffusion and their analysis are challenging problems due to the non-smooth nature of the Brownian motion (M¨orters & Peres, 2010).
There has been a signiﬁcant progress being made however in devising and analyzing Langevin MCMC algorithms, e.g. Cheng et al.
(2018) introduced an algorithm with a mixing time that scales as O ( √ d ) in a notable contrast to the best known O ( d ) for MCMC algorithms based on overdamped Langevin diffusion.
The dimension dependence of the mixing time is of great interest here since we expand the dimension M-fold.
To give an idea regarding the dimension, for the 4 ⊗ 8 model on FFHQ-256,
Md ≈ 10 6 .
We implemented (Cheng et al., 2018,
Algorithm 1) in this paper which to our knowledge is its ﬁrst use for generative modeling.
In addition, we used a Langevin MCMC algorithm due to Sachs et al.
(2017), also used by Arbel et al.
Note, in addition to γ and u , Langevin MCMC requires δ , the step size used for discretizing Eq. 23.
[Figure]
[Figure]
We used Algorithm 1 with δ = 2 , γ = 1 / 2 , u = 1 .
Transitions are best viewed electronically.
E XPERIMENTAL R ESULTS
For all models with the MDAE and MEM 2 parametrizations, we used U 2 -Net
(Qin et al., 2020) network architecture, a recent variant of UNet (Ronneberger et al., 2015), with a few simple changes (see Appendix E).
For the MUVB parametrization for MNIST (evaluated in Sec.
G.3), we used
Residual networks for the encoder, decoder and metaencoder.
Despite the inherent complexity of the task of learning high-dimensional energy models, it is notable that our framework results in a single least-squares objective function, which we optimized using Adam (Kingma & Ba, 2014) without additional learning rate schedules.
Additional details of the experimental setup are in Appendix E .
Arguably, the ultimate test for an energy model is whether one can generate realistic samples from it with a single fast mixing MCMC chain that explores all modes of the distribution indeﬁnitely , starting from random noise.
We informally refer to them as lifelong Markov chains .
To give a physical picture, a gas of molecules that has come in contact with a thermal reservoir does not stop mid air after thermalizing—arriving at its “ﬁrst sample”—it continues generating samples from the Boltzmann distribution as long as the physical system exists.
To meet this challenge, the energy landscape must not have any pathologies that cause the sampling chain to break or get stuck in certain modes.
In addition, we need fast mixing (Langevin) MCMC algorithms, a very active area of research by itself.
To put this goal in context, in recent energy models for high-dimensional data (Nijkamp et al., 2019;
Du & Mordatch, 2019; Du et al., 2020), sampling using MCMC quickly breaks or collapses to a mode (chains longer than a few hundred steps have never been reported).
Thus, evaluation in prior work relies on samples from independent MCMC chains.
We also note that denoising models such as denoising diffusion models (Gao et al., 2020) produce samples of much higher visual quality currently, but in those approaches the chain is conditional and it stops after generating a clean sample.
Our work was motivated with this goal in mind.
Since the main text is devoted to the theoretical framework, we summarize the main experimental results here and refer to the appendices for details.
For MNIST, we obtain fast-mixing chains for over 1 M steps using MDAE and MUVB.
On CIFAR10 and FFHQ-256, we obtain stable chains up to 1 M and 10 K steps, respectively, using MDAE.
The results are remarkable since energy models—that learn a single energy/score function—have never been scaled to 256 × 256 resolution images.
The closest results are by Nijkamp et al.
and Du et al.
(2020) on CelebA (128 × 128)—in particular, our results in Figs.
3, 4 can be compared to (Nijkamp et al., 2019, Figs. 1, 2): they also visualize unconditional fast mixing chains.
The chains mix well for MNIST/FFHQ-256, but for CIFAR-10, we observe that the MDAE chain cycles through certain modes, leading to a poor FID score of 79 despite lacking image artifacts (Figs. 11, 12).
Our detailed results are organized in appendices as follows.
Appendix F is devoted to an introduction to Langevin MCMC with demonstrations using our FFHQ-256 4 ⊗ 8 model.
Appendix G is devoted to MNIST experiments.
At ﬁrst we compare 1 / 4 ⊗ 1 and 1 ⊗ 16 (Sec. G.1).
These models are statistically identical regarding the plug-in estimators (Eq. 10) and they arrive at similar training losses, but very different as generative models.
This sheds light on the question why we considered such high noise levels in designing experiments (Fig. 5).
We then present the effect of increasing M for a ﬁxed σ , the issue of time complexity, and mixing time vs. image quality trade-off (Sec.
G.2).
The discussions in Sec. G.1 and Sec.
G.2 are closely related.
We then compare MDAE and MUVB for 1 ⊗ 4 with the message that MDAE generates sharper samples while MUVB has better mixing properties (Sec. G.3).
We then present one example of a lifelong Markov chain (Sec. G.4).
Our
MDAE model fails at CIFAR-10 challenge due to the optimization problems we have encountered.
The results and further discussion are presented in Appendix H .
Finally, in Appendix I we show chains obtained for FFHQ-256, including the ones that some of the images in Fig.
4 were taken from.
R ELATED W ORK
Despite being a fundamental framework for denoising, empirical Bayes has been surprisingly absent in the DAE/DSM literature.
Raphan & Simoncelli (2011) ﬁrst mentioned the connection between empirical Bayes and score matching, but this work is practically unknown in the literature; for much of its development, DSM was all tied to DAEs (Vincent, 2011; Alain & Bengio, 2014); see (Bengio et al., 2013) for an extensive survey.
Abstracting DSM away from DAE is due to Saremi et al.
(2018) where they directly parametrized the energy function.
In this work, we closed the circle, from empirical Bayes to DAEs .
One can indeed use the MDAE objective and learn a generative model: a highly generalized DAE naturally arise from empirical Bayes.
The important subtlety here is, as we remarked in Sec. 3.2, we can only arrive at this connection starting from our generalization of empirical Bayes, not the other way around.
Of course, this core multimeasurement aspect of our approach, without which we do not have a generative model, does not exist in the DAE literature.
To address the problem of choosing a noise level in DSM (Saremi et al., 2018), Song & Ermon (2019) studied it with multiple noise levels by summing up the losses using a weighing scheme.
See (Li et al., 2019; Chen et al., 2020; Song & Ermon, 2020; Kadkhodaie & Simoncelli, 2020; Jolicoeur-Martineau et al., 2020) in that direction.
The learning objectives in these models are based on heuristics in how different noise levels are weighted.
Denoising diffusion models (Ho et al., 2020; Song et al., 2020;
Gao et al., 2020) follow the same philosophy while being theoretically sound.
Sampling in these models are based on annealing or reversing a diffusion process.
Our philosophy here is fundamentally different.
Even “noise levels” have very different meaning here, associated with the M noise channels of the factorial kernel.
All noise levels (which we took to be equal in later parts, a meaningless choice in other methods) are encapsulated in a single energy/score function.
Using Langevin MCMC we only sample highly noisy data y .
Clean data in R d is generated via a single step by computing (cid:98) x
( y ) .
C ONCLUSION
We approached the problem of generative modeling by mapping a complex density in R d to a smoother “dual” density, named M-density, in R Md .
Permutation invariance is a design choice which we found particularly appealing.
Using factorial kernels for density estimation and generative modeling have never been explored before and this work should open up many new avenues.
We believe the topic of parametrization will play an especially important role in future developments.
There is a unity in the parametrization schemes we proposed and studied in the paper, but much remains to be explored in understanding and extending the relationship between MEMs and MSMs.
Our MDAE parametrization scheme (an instance of MSM) is especially appealing for its simplicity and connections to the past literature.
DAEs have a very rich history in the ﬁeld of representation learning.
But research on them as generative models stopped around 2015.
Our hypothesis is that one must use very large noise to make them work as generative models for complex datasets in high dimensions.
Our permutation invariant multimeasurement approach is an elegant solution since it allows for that choice at the cost of computation (large M ), with only one parameter left to tune: σ !
Crucially, the probabilistic interpretation is ingrained here, given by the algebraic relations between the MDAE output ν ( y ) , the Bayes estimator (cid:98) x ( y ) , and the score function ∇ log p ( y ) .
and Eq.
18 here.
However, one is on barely noisy data ( σ → 0 ) in R d , the other on multimeasurements in R Md for any (high) noise level σ .
There, they arrived at Eq.
5 (with some effort) starting from a DAE objective.
Here, we start with Eq.
18 (it deﬁnes the score function) and arrive at MDAE objective (Eq. 19), in one line of algebra.
E THICAL C ONSIDERATIONS
By their very nature, generative models assign a high (implicit or explicit) likelihood to points in their training set.
When samples are generated from a trained generative model, it is quite possible that some samples are very similar or identical to certain data points in the training set.
This is important to keep in mind when the data used to train a generative model is conﬁdential or contains personally identiﬁable information such as pictures of faces in the FFHQ-256 dataset (Karras et al., 2019).
Care should be taken to abide by license terms and ethical principles when using samples from generative models.
For FFHQ-256 in particular, please see terms of use at https: //github.com/NVlabs/ffhq-dataset/blob/master/README.md .
Additionally, we recommend considering the in-depth discussion on this subject by Prabhu & Birhane (2020) before deploying applications based on generative models
R EPRODUCIBILITY S TATEMENT
All important details of the experimental setup are provided in
Appendix E. Our code is currently proprietary but will be made available to reviewers conﬁdentially if requested.
We are also working on making our code public upon acceptance of the paper
A G AUSSIAN M- DENSITIES
We study the Bayes estimator for Gaussian MNMs analytically for X = N ( µ, σ 20 I d ) , where µ ∈ R d .
As in the example for Poisson MNM in the paper, the ﬁrst step is to derive an expression for the joint density
[Equation]
We start with M = 2 .
We used Mathematica to perform the integral above (it is a surprisingly cumbersome integral):
[Equation]
where Z is the partition function: Z ( σ 0 , σ 1 , σ 2 )
=
[Equation]
2 π
[Equation] [Equation]
2 ( σ − 2 0 + σ − 2 1 + σ − 2 2 ) (cid:1) d/ 2 .
The multimeasurement Bayes estimator of X is computed next via Eq.
[Equation]
As a sanity check, (cid:98) x ( y ) is the same whether Eq.
7 is computed using m = 1 or m = 2 .
Another sanity check is the limit σ 1 → 0 , which we recover (cid:98) x
( y 1 , y 2 ) =
y
= x .
The linear relation (linear in y for µ = 0 ) observed in Eq.
26 is a generalization of the well known results for M = 1 , known as Wiener ﬁlter (Wiener, 1964).
A visualization of this result for a mixture of Gaussian is given in Fig.
Performing the integralis cumbersome, but there is a pattern that emerges (in our implementation in Mathematica) for the case M > 2 which we checked for M = 3 and M = 4 .
The results are stated below, and then we derive the expression for the Bayes estimator (cid:98) x ( y ) .
It is convenient to deﬁne the following notations:
[Equation]
y M \ m , where y M \ m :=
y m (cid:89) l ∈ [ M ] \ m σ 2 l , and y 0 :
= µ
[Equation]
With these notations, the expression for the M-density is given by:
[Equation]
We can now derive (cid:98) x
( y ) using Eq.
Given the expression above, the algebra is straightforward:
[Equation]
Note that the ﬁnal expression is the same by carrying the algebra for any measurement index m .
B M- DENSITY SCORE MATCHING ≡ M ULTIMEASUREMENT DSM
Here we provide the score matching formalism (Hyv¨arinen, 2005) on M-densities and we arrive at a multimeasurement denoising score matching (MDSM) learning objective.
The derivation closely follows (Vincent, 2011) using our notation.
In what follows equalities are modulo additive constants (not functions of parameters θ ) which we denote by the color gray:
[Equation]
The M-density score matching learning objective is given by the ﬁrst identity above and the last equality is the MDSM objective:
[Equation]
For Gaussian MNMs
[Equation]
the MDSM learning objective simpliﬁes to:
[Equation]
C MUVB: AN UNNORMALIZED LATENT VARIABLE MODEL
In this section we give a formulation of the energy function parametrization scheme which is named multimeasurement unnormalized variational Bayes (MUVB).
Conceptually, MUVB is motivated by the goal of connecting empirical Bayes’ denoising framework with variational Bayes’ inference machinery.
There seem to be fundamental challenges here since one model is based on leastsquares estimation, the other based on the principle of maximum likelihood.
In what follows, this “fundamental conﬂict” takes the shape of latent variable model that is unnormalized .
There have been recent sudies along these lines (Saremi, 2020a;b) that we expand on; in particular we introduce the novel metaencoder .
At a high level, the idea here is to set up a latent variable model for M-densities and use the variational free energy as the energy function.
We refer to (Jordan et al., 1999) for an introduction to variational methods.
In what follows, we use the lighter notation for σ ⊗ M models (Remark 2) that we also used in Sec. 4.
In its simplest form, the latent variable model for M-densities is set up by the following choices:
• The prior over latent variables which we take to be Gaussian p
( z ) =
N ( z ; 0 , I d z ) .
• The approximate posterior over latent variables q φ ( z | y ) which is taken to be the factorized Gaussian, a standard choice in the literature.
• For the conditional density
p η
( y | z ) we can consider the following p η ( y | z )
=
N ( y ; ν η ( z ) , σ 2 I Md ) ,
which may seem as a especially “natural” choice since samples from the M-density are obtained by adding multimeasurement noise to clean samples.
In the variational autoencoder parlance, φ are the encoder ’s parameters, η the decoder ’s parameters, and
ν = ( ν 1 , . .
. , ν M ) the decoder’s outputs.
With this setup, the variational free energy is easily derived which we take to be the energy function as follows ( θ = ( φ, η ) ):
[Equation]
where the KL divergence is derived in closed form, and the expectation is computed by sampling z ∼ q φ ( z | y ) via the reparametrization trick (Kingma & Welling, 2013).
In this machinery, in the inference step we are forced to have normalized (and tractable) posteriors since we must take samples from q φ ( z | y ) .
What about the conditional density p ( y | z ) ?
If we were to formalize a VAE for modeling M-densities we had to have a normalized conditional density p ( y | z ) since that framework is based on the principle of maximum likelihood .
What is intriguing about our setup is that since our goal is to learn the unnormalized p ( y ) , we can consider conditional densities p ( y | z ) that are unnormalized .
An intuitive choice is the following that we name metalikelihood :
[Equation]
Now, the energy function takes the following form ( θ = ( φ, η, ζ ) ):
[Equation]
This is the ﬁnal form of MUVB energy function.
We refer to h ζ by metaencoder.
Its input is ( y , ν
( z ; η ))
(note that z itself is a function of y via inference), and one can simply use any encoder architecture to parametrize it.
In our implementation, we used the standard VAE pipeline and we reused the encoder architecture to parametrize the metaencoder.
D T WO W ALK -J UMP S AMPLING A LGORITHMS
In this section we provide the walk-jump sampling (WJS) algorithm based on different discretizations of underdamped Langevin diffusion (Eq. 23).
The ﬁrst is due to Sachs et al.
(2017) which has also been used by Arbel et al.
(2020) for generative modeling.
The second is based on a recent algorithm analyzed by Cheng et al.
(2018) which we implemented and give the detailed algorithm here.
In addition, we extended their calculation to general friction (they set γ = 2 in the paper) and provide an extension of the (Cheng et al., 2018, Lemma 11) for completeness (the proof closely follows the calculation in the reference).
In both cases we provide the algorithm for σ ⊗ M models.
D. W ALK -J UMP S AMPLING A LGORITHM I
( y ) or score function g ( y )
≈ ∇ log p ( y ) 3: Ouput (cid:98) X K 4: Y 0 ∼ Unif ([0 , 1] Md ) 5: V
←6: for k = [0 , . . .
, K ) do 7:
Y k +1 ← Y k
+ δ V
k
Ψ k +1 ← −∇ y f
( Y k +1 ) or Ψ k +1 ← g
( Y k +1 ) 9:
V k +1 ← V k
+
uδ Ψ k
+1 /10:
B k +1 ∼ N (0 , I Md ) 11:
V k +1 ← exp( − γδ )
V k
+
 +
uδ Ψ k +1 /
+ (cid:112) u
(1 − exp( − 2 γδ ))
B k + 12:
Y k +1 ← Y k
+
+ δ V
k +1 /
: end for 14:
(cid:98) X K ← (cid:104) Y K,m − σ 2 ∇ m f ( Y K ) (cid:105) or (cid:98) X K ← (cid:104) Y K,m + σ 2 g m ( Y K ) (cid:105)
D. W ALK -J UMP S AMPLING A LGORITHM II
Before stating the algorithm, we extend the calculation in Cheng et al.
(2018, Lemma 11) to general friction, which we used in our implementation of their Langevin MCMC algorithm.
The algorithm is stated following the proof of the lemma which closely follows Cheng et al.
The starting point is solving the diffusion Eq. 23.
With the initial condition ( y 0 , v 0 ) , the solution ( y t , v t ) to the discrete underdamped Langevin diffusion is ( t is considered small here)
[Equation]
These equations above are then used in the proof of the following lemma.
Lemma 1.
Conditioned on ( y 0 , v 0 ) , the solution of underdamped Langevin diffusion (Eq. 23) integrated up to time t is a Gaussian with conditional means
[Equation]
and with conditional covariances
[Equation]
Proof.
Computation of conditional means is straightforward as the Brownian motion has zero means:
[Equation]
All the conditional covariances only involve the Brownian motion terms.
We start with the simplest:
[Equation]
From Eq.
[Equation]
the Brownian motion term for y t is given by
[Equation]
The conditional covariance for y t follows similar to Eq.
[Equation]
Finally the conditional covariance between y t and v t is given by
[Equation]
Lemma 1.
Below, we provide the Cholesky decomposition of the conditional covariance matrix.
−∇
V k +1 ← V k e − γδ + γ − 1 u
(1 − e − γδ )
Ψ k
Y k +1 ← Y k + γ − 1 (1 − e − γδ )
V k + γ − 1 u (cid:0)
δ − γ − 1 (1 − e − γδ ) (cid:1)
Ψ k
Md
− σ
∇ f
The Cholesky decomposition Σ = LL (cid:62) for the conditional covariance in Lemma 1 is given by:
[Equation]
Md
[Equation]
Md
[Equation]
38) where
[Equation]
Algorithm 1 vs. Algorithm
The two algorithms presented in this section have very different properties in our high dimensional experiments.
In general we found Algorithm 1 (A1) to behave similarly to overdamped version (see Appendix F for some illustrations) albeit with much faster mixing.
We found friction γ to play a very different role in comparing the two algorithms, also apparent in comparing the mathematical expressions in both algorithms: in A1, the friction only appears in the form exp( − γδ ) ; it is more complex in A2.
Similar discrepancy is found regarding u .
Remark 4 (Initialization scheme) .
The initialization Y 0 ∼ Unif ([0 , 1] Md ) used in the algorithms is the conventional choice, but for very large σ we found it effective to initialize the chain by further adding the Gaussian noise in R Md :
[Equation]
Otherwise, with large step sizes δ
= O ( σ ) the chains starting with the uniform distribution could break early on.
We found this scheme reliable, and it is well motivated since M-density p ( y ) is obtained by convolving p ( x ) with the Gaussian MNM.
With this initialization one starts “relatively close” to the M-density manifold and this is more pronounced for larger σ .
However, this initialization scheme is based on heuristics (as is the conventional one).
Remark 5 (The jump in WJS is asynchronous) .
The jump in WJS is asynchronous (Fig. 2) and it can also be taken inside the for loop in the algorithms provided without affecting the Langevin chain.
Remark 6 (The ∆ k notation) .
In running long chains we set K to be large, and we report the jumps (taken inside the for loops in the algorithms provided) with a certain ﬁxed frequency denoted by ∆ k .
Remark 7 (WJS for Poisson M-densities) .
WJS is a general sampling algorithm, schematized in
Fig.
However, for Poisson MNMs it will involve sampling a discrete distribution in N Md which we do not know how to do efﬁciently in high dimensions (note that Langevin MCMC cannot be used since the score function is not well deﬁned for discrete distributions).
We can use MetropolisHastings algorithm and its popular variant Gibbs sampling but these algorithms are very slow in high dimensions since at each iteration of the sampling algorithm in principle only 1 dimension out of Md shall be updated, so in the general case the mixing time is at best of O ( Md ) ; one can use “blocking” techniques but they are problem dependent and typically complex (Andrieu et al., 2003).
E E XPERIMENTAL S ETUP
Datasets Our experiments were conducted on the MNIST (LeCun et al., 1998),
CIFAR-
(Krizhevsky, 2009) and FFHQ-256 (Karras et al., 2019; pre-processed version by Child, 2020) datasets.
The CIFAR-10 and FFHQ-256 training sets were augmented with random horizontal ﬂips during training.
Network Architectures For all experiments with the MDAE and MEM 2 parametrization, the
U 2 Net network architecture (Qin et al., 2020) was chosen since it is a recent and effective variant of the UNet (Ronneberger et al., 2015) that has been successfully applied to various pixel-level prediction problems.
However, we have not experimented with alternatives.
Minor adjustments were made to the U 2 Net from the original paper: we removed all normalization layers, and switched the activation function to x (cid:55)→ x · sigmoid ( x ) (Elfwing et al., 2017; Ramachandran et al., 2017).
To approximately control the variance of inputs to the network in each dimension, even if the noise level σ is large, the noisy input values were divided by a scaling factor √ 0 .
225 2 + σ 2 .
We also adjusted the number of “stages” in the network, stage heights and layer widths (see Table 1 for details), in order to ﬁt reasonable batch sizes on available GPU hardware and keep experiment durations limited to a few days (see Table 2).
For experiments with the MUVB parametrization (the MNIST model in Sec.
G.3), we used Residual convolutional networks utilizing a bottleneck block design with skip connections and average pooling layers (Lin et al., 2013;
Srivastava et al., 2015; He et al., 2016).
The architecture of the encoder and metaencoder were kept the same, while the decoder used the encoder architecture in reverse, with the pooling layers substituted with nearest neighbour upsampling.
See Table 3 for a detailed description of the encoder architecture.
Remark 8.
Strictly speaking, for σ ⊗ M models the energy function is permutation invariant and y should be treated as { y 1 , . . .
, y M } , a set consisting of M Euclidean vectors.
This is however difﬁcult to enforce while keeping the model expressive .
We note that in our model the permutation symmetry is a “fragile symmetry” which one can easily break by making σ m just slightly different from each other.
Therefore, we designed the architecture to be used for the general setting of σ m different from each other, and in experiments we just set them to be equal.
Training All models were trained using the Adam (Kingma & Ba, 2014) optimizer with the default setting of β 1 = 0 .
9 , β 2 = 0 .
999 and (cid:15) =
10 − 8 .
Table 2 lists the main hyperparameters used and hardware requirements for training MDAE models for each dataset.
The resulting training curves are shown in Fig.
6, showing the stability of the optimization with a relatively simple setup.
Software All models were implemented in the CPython (v3.8) library PyTorch v1.9.
(Paszke et al., 2017) using the PyTorch Lightning framework v1.4.6 (Falcon et al.).
[Table]
[Figure]
[Table] [Table]
F L ANGEVIN MCMC L ABORATORY
In this section, we give a light tutorial on Langevin MCMC used in the walk phase of WJS.
Perhaps the most interesting result enabled by multimeasurement generative models is that Langevin MCMC in high dimensions becomes stable enough for us to examine and dissect the performance of various
Langevin MCMC variants in practice.
The reader might ask why we used underdamped Langevin MCMC versus its well-known overdamped version.
The reason is the recent theoretical developments discussed in Sec. 5 that show much faster dimension dependence for the convergence (mixing time), O ( √ d ) vs. O ( d ) , for underdamped
Langevin MCMC.
To give an idea regarding the dimension, for the 4 ⊗ 8 model on FFHQ-256,
Md ≈ 10 6 .
Our trained model for 4 ⊗ 8 is used throughout this section.
In addition, the MCMC is initialized with the ﬁxed random seed throughout this section.
F. O VERDAMPED L ANGEVIN MCMC
We start with overdamped Langevin MCMC which is simpler and give us some intuitions on the behavior of Langevin MCMC in general.
The algorithm is based on the following stochastic iterations (setting inverse mass u = 1 ):
[Equation]
B
[Equation]
, where ε k ∼ N (0 , I Md ) .
Comparing above with Eq.
7, it is intuitive to consider setting δ
= O ( σ ) .
We start with δ = σ : in this case A k pulls noisy data to clean manifold (note the important extra factor of 1 / 2 however) and B k term corresponds to adding multimeasurement noise sampled from N (0 , σ 2 I Md ) .
The overall effect is that one starts “walking” on the manifold of the M-density.
And if that is the case, the jumps computed by (cid:98) x ( y k ) should give samples close to the manifold of the (clean) data distribution—samples from p ( x ) .
We start with that experiment by setting δ = σ .
Note that in walk-jump sampling we sample the M-density in R Md generating highly noisy data for σ = 4 ; we only show clean (jump) samples in R d .
This is shown next by skipping every 30 steps, i.e. ∆ k = 30 :
Our intuition was partly correct in that the sampler manages to “capture modes” but the step size is clearly too high.
Now consider δ = σ/ 2 and ∆ k = 100 :
Continuing the chain, now shown for ∆ k = 1000 , we arrive at
Theoretically, to converge to the true distribution (associated with the M-density) the step size should be small, where “small” is dictated by how “close” one wishes to be to the true distribution, e.g. see (Cheng et al., 2018, Theorem 1).
As we see above, even for δ = 2 the sampler does not converge after 1000 steps, but a ﬁxed δ = σ/ 2 appears to be too high for this model.
By continuing with the chain with the step size δ = σ/ 2 one starts to gradually move away from the manifold as shown next:
F. U NDERDAMPED L ANGEVIN MCMC
In this section we focus on underdamped Langevin MCMC.
The algorithms are based on discretizing underdamped Langevin diffusion (Eq. 23).
In this paper, we considered two such discretization schemes, by Sachs et al.
(2017) that we used in Algorithm 1, and the one due to Cheng et al.
(2018) that we implemented and used in Algorithm 2.
The ﬁrst algorithm is not analyzed but it is very easy to show that in the limit γ → ∞ it will be reduced to the overdamped Langevin MCMC we just discussed.
Second algorithm is more complex in how friction γ and inverse mass u behave.
For comparison, we present results obtained by both algorithms, showing Algorithm 1 ﬁrst.
As in the previous section, we start with δ =
σ ; we set γ = 1 ( u = 1 unless stated otherwise) with ∆ k = 30 :
As in the previous section δ =
σ is too high for FFHQ-256, 4 ⊗ 8 model.
(However, note the relative stability of Algorithm 2, for this random seed.)
Next, we consider δ = σ/ 2 , γ = 1 , shown with ∆ k = 100 :
Note the remarkable contrast between the results above and the corresponding ﬁgure (for ∆ k = 100 ) in the previous section.
We continue the chain with the same ∆ k
= 100 :
Both algorithms start to diverge for this choice of parameters ( δ = σ/ 2 , γ = 1 ).
To add more stability we can consider increasing the “mass” (lowering u ).
In (Cheng et al., 2018), u is replaced with 1 /L , where L is the Lipschitz constant of the score function.
By lowering u , we are telling the algorithm that the score function is less smooth and MCMC uses more conservative updates.
This can also be seen in the equations in both algorithms, where the step size δ is multiplied by u in several places.
However, note this effect is more subtle than just simply scaling δ , especially so for Algorithm 2 as can be seen by inspection.
Next, we consider u = 1 / 2 .
This will affect the mixing time and the chains are now shown with ∆ k = 1000 (instead of ∆ k = 100 earlier):
This experiment concludes this section.
We come back to more FFHQ-256 chains in Appendix I.
G MNIST G. M ORE IS D IFFERENT
There is a notion of congruence which arise in our analysis in Sec.
2.3 which we denote by “ ∼ = ”: two different models, σ ⊗ M and σ (cid:48) ⊗ M (cid:48) , “agree with each other” in terms of the plug-in estimator
[Equation]
However these models are vastly different, one is in R Md , the other in R M (cid:48) d .
And if σ (cid:29) σ (cid:48) we expect that σ ⊗ M model to have better properties as a generative model.
This is clear in the regime ( σ (cid:48) (cid:28) 1 , M (cid:48) = 1) .
In that regime, almost no learning occurs: at test time MCMC either gets stuck in a mode for a long time (the best case scenario) or simply breaks down.
We refer to (Saremi &
Hyv¨arinen, 2019) for a geometric intuition on why large noise is essential in these models regarding the walk-jump sampling.
We should point out that this is a general phenomenon in all denoising objectives (Alain & Bengio, 2014; Song & Ermon, 2020).
In a nutshell, when noise is small
the
Langevin sampler ﬁnds it difﬁcult to navigate in high dimensions and the forces of the Brownian motion eventually takes it to a part of space unknown to our model and MCMC breaks down (the ﬁrst experiment below).
If we are “lucky” we can ﬁnd the manifold initially (the second experiment below) and MCMC may remain stable for a while but a model with σ > σ (cid:48) will have better mixing properties.
We should highlight that in the analysis of Cheng et al.
the mixing time scales as κ 2 where κ is the condition number (see Sec. 1.4.1) proportional to the Lipschitz constant; this quantiﬁes the intuition that MCMC on smoother densities (smaller Lipschitz constant) mixes faster.
To validate these intuitions we consider comparing 1 ⊗ 16 and 1 / 4 ⊗ 1 .
In addition to statistically identical plug-in estimators, these models also arrive at similar training losses with the same training schedules (1.03 for 1 ⊗ 16 and 1.19 for 1 / 4 ⊗ 1 ).
We start with 1 / 4 ⊗ 1 .
Below we show WJS using
Algorithm 2 with the setting ( δ = 1 / 8 , γ = 1 / 2 , u = 1) (the setting δ = σ/ 2 , γ = 1 / 2 , u = 1 is a simple choice we used in testing all our σ ⊗ M models).
Here, MCMC breaks down shortly after the last step shown here (this chain is shown with ∆ k = 30 , in the remainder ∆ k = 500 ):
Algorithm 1 is more stable here and we arrive at the following chain (total of 10 5 steps):
Now, for 1 ⊗ 16 model, using Algorithm 2 with parameters ( δ = 1 / 2 , γ = 1 / 2 , u = 1) , we arrive at:
And using Algorithm 1 with the same parameters (and initial seed) we arrive at: (Note the stark differences in terms of mixing between Algorithms 1 and 2 for this model.)
In summary, more is different : quantitative differences become qualitative ones (Anderson, 1972) and this is highlighted here for “congruent” σ ⊗ M models that differ in terms of the noise level.
If computation is not an issue, one should always consider models with larger σ but this remains a conjecture here in its limiting case (as σ and M grow unbounded, while σ/ √ M remains “small”).
G. I NCREASING M : TIME COMPLEXITY , SAMPLE QUALITY , AND MIXING TIME TRADE - OFFS
Here we extend the ablation study in the previous section by studying σ ⊗ M models with a ﬁxed σ where we increase M .
In the previous section we studied two models where (cid:98) x
( y ) has the same statistical properties as measured by the loss (and the plug-in estimator), and we validated the hypothesis that MCMC sampling has better mixing properties for the smoother M-density (with larger σ ).
Now, the question is what if we keep σ ﬁxed and increase M ?
As we motivated in Sec. 1, by increasing M the posterior p ( x | y ) concentrates on the mean (cid:98) x ( y ) , thus increasing the sample quality in WJS.
However there are trade-offs: (i)
The ﬁrst is the issue of time complexity which is an open problem here.
In our architecture, the M measurements are passed through a convolutional layer and after that models with different M have approximately the same time complexity.
The problem is this may not be an ideal architecture.
Presumably, one should increase the network capacity when increasing M (for the same σ ) but by how much we do not know (this clearly depends on how complex p X is).
Here we ignore this issue and we assume that our architecture has enough capacity for the largest M = 16 considered.
(ii)
The second trade-off is as we increase M , the price we pay for higher sample quality is that we will have longer mixing times, both in terms of generating the ﬁrst sample, and more importantly in the mixing between modes.
This trade-off is intuitive since p ( y 1 , . . .
, y M ) becomes more complex for larger M (for ﬁxed σ ).
Below we compare WJS chains (as before, without warmup) for σ = 1 and M = 1 , 2 , 4 , 16 :
We used Algorithm 2 ( δ = 1 / 2 , γ = 1 / 2 , u = 1) with the same initial seed.
Note that, at the cost of sample quality, all classes are visited in 1 ⊗ 1 in the short chain presented.
G. MUVB VS . MDAE
Below we make a one to one comparison between MUVB and MDAE for the 1 ⊗ 4 setting we trained on MNIST.
In both cases we ran a single chain for 1 million+ steps.
For sampling M-density we used the Langevin MCMC algorithm by Sachs et al.
(2017) with parameters ( δ = 1 , γ = 1 / 4 , u = 1) (see
Appendix D).
This is an “aggressive” choice of parameters for Langevin MCMC, designed for fast mixing, yet the chains do not break up to 1M+ steps (we stopped them due to computational reasons).
Below we show the two chains at discrete time resolution of 5 steps ( the ﬁrst sample, top left corner, is obtained after only 5 steps ).
This comparison visually demonstrates that MUVB has better mixing and MDAE better sample quality.
All digit classes are visited in MUVB in a variety of styles insteps shown here.
MDAE chain is cleaner but the classes { 0 , 1 , 5 , 6 } are visited scarcely.
[Figure]
We demonstrate the WJS chain with 1 million+ steps obtained in MDAE ( 1 ⊗ 4 ) in its entirety (viewed left to right, top to bottom).
We informally refer to these long chains that never break as lifelong
Markov chains .
The chain is reported here in its “entirety” to demonstrate the fact that it does not get stuck in a mode and remains stable through the end (we stopped it due to computational reasons).
H CIFAR-10: MDAE FAILURE MODE
The main “fear” in parametrizing the score function directly in MSM, which MDAE is an instance of, is we are not guaranteed to learn a conservative score function (see Sec. 4.1).
More importantly, we do not have a control over how the learned model is failing in that regard.
These fears were realized in our
CIFAR experiments on the 1 ⊗ 8 model presented below.
The main results use Algorithm 2.
Algorithm 1 simply fails for the Langevin MCMC parameters we experimented with below (for completeness we present a short chain at the bottom panel).
In 2D toy experiments the non-conservative score functions were quantiﬁed in (Saremi et al., 2018, Fig. 1h), but this quantiﬁcation is difﬁcult in high dimensions.
However, what is intriguing about the experiments here is that the “non-conservative nature” of the score function is manifested—this is merely a conjecture—in the cyclic mode visits apparent in the Markov chain.
We encountered the same problem using spectral normalization (Miyato et al., 2018): the chain below was uptained using Algorithm 2 with the setting ( δ = 1 / 2 , γ = 1 / 2 , u = 1) where we got the
FID score of 99.
We ran the chain for 1 M steps; the ﬁrst 600 K steps are visualized below:
I FFHQ-
In this section, we provide WJS chains for three models we trained on FFHQ-256:
MDAE 4 ⊗ 8 ,
MDAE 4 ⊗ 4 , and MEM 2 4 ⊗ 4 (without metaencoder).
We refer to Algorithm 1 by A1 and
Algorithm 2 by A2 , and MCMC parameters are listed as ( δ, γ, u ) .
I. MDAE
[Table]
A1; (2 , 1 , 1) ; ∆ k =
A2; (2 , 1 , 1) ; ∆ k =
A2; ( 1 / 10 , 1 , 10) ; ∆ k =
A2; (4 , 1 / 2 , 1 / 4 ) ; ∆ k =
A1; (2 , 1 / 2 , 1) ; ∆ k = A1; (2 , 1 / 2 , 1) ; ∆ k =
A2; ( 1 / 10 , 2 , 50) ; ∆ k =
A1; (2 , 1 / 2 , 1) ; ∆ k = A1; ( 3 / 2 , 1 / 2 , 1) ; ∆ k =
I.
[Figure]
⊗
[Figure] [Figure]