Approximate Probabilistic Inference with Composed Flows
Authors: Jay Whang
Erik Lindgren
Alex Dimakis
Venue:  (2021)
Keywords: normalizing flow
probabilistic inference
variational inference
inverse problem

Abstract: We study the problem of probabilistic inference on the joint distribution defined by a normalizing flow model. Given a pre-trained flow model $p(\boldsymbol{x})$, we wish to estimate $p(\boldsymbol{x}_2 \mid \boldsymbol{x}_1)$ for some arbitrary partitioning of the variables $\boldsymbol{x} = (\boldsymbol{x}_1, \boldsymbol{x}_2)$. We first show that this task is computationally hard for a large class of flow models.  Motivated by this hardness result, we propose a framework for $\textit{approximate}$ probabilistic inference. Specifically, our method trains a new generative model with the property that its composition with the given model approximates the target conditional distribution. By parametrizing this new distribution as another flow model, we can efficiently train it using variational inference and also handle conditioning under arbitrary differentiable transformations. Since the resulting approximate posterior remains a flow, it offers exact likelihood evaluation, inversion, and efficient sampling. We provide an extensive empirical evidence showcasing the flexibility of our method on a variety of inference tasks with applications to inverse problems. We also experimentally demonstrate that our approach is comparable to simple MCMC baselines in terms of sample quality. Further, we explain the failure of naively applying variational inference and show that our method does not suffer from the same issue.

Under review as a conference paper at ICLR 2021
APPROXIMATE PROBABILISTIC INFERENCE WITH COMPOSED FLOWS
Anonymous authors Paper under double-blind review
ABSTRACT
We study the problem of probabilistic inference on the joint distribution defined by a normalizing flow model. Given a pre-trained flow model p(æ). we wish to estimate p(a²2 a1) for some arbitrary partitioning of the variables We first show that this task is computationally hard for a large class of flow mod- els. Motivated by this hardness result, we propose a framework for approximate probabilistic inference. Specifically, our method trains a new generative model with the property that its composition with the given model approximates the target conditional distribution. By parametrizing this new distribution as another flow model, we can efficiently train using variational inference and also handle conditioning under arbitrary differentiable transformations. Since the resulting ap- proximate posterior remains flow, it offers exact likelihood evaluation, inversion, and efficient sampling. We provide extensive empirical evidence showcasing the flexibility of our method on a variety of inference tasks with applications to inverse problems. We also experimentally demonstrate that approach is comparable to simple MCMC baselines in terms of sample quality. Further, we explain the failure of naively applying variational inference and show that our method does not suffer from the same issue.
(æ1, 222)
INTRODUCTION
Generative modeling has seen unprecedented growth in the recent years. Building on the success of deep learning, deep generative models have shown impressive ability to model complex distributions in variety of domains and modalities. Among them, normalizing flow models (see Paaaaakarios ett al. (2019) and references therein) stand out due to their computational flexibility, as they offer efficient sampling. likelihood evaluation, and inversion. While other types of models currently outperform flow models in terms of likelihood and sample quality, flow models have the advantage that they are relatively easy to train using maximum likelihood and do not suffer from issues that other models possess (e.g. mode collapse for GANs, posterior collapse for VAEs, slow sampling for autoregressive models). These characteristics make nornalizing flows attractive for variety of downstream tasks, including density estimation, inverse problems, superviee learning. reinforcement learning, and audio synthesis (Ho al. 2019 Asim et al.
2019 Atanov et al. 2019
Ward et al.) 2019; Oord al. 2018).
Even with such computational flexibility, how to perform efficient probabilistic inference on a flow model still remains largely unknown. This question is becoming increasingly important as generative models increase in size and the computational resources necessary to train them from scratch are out of reach for many researchers and practitioner:" If it was possible to perform probabilistic inference on flow models, we could re-purpose these powerful pre-trained generators for numerous custom tasks.
This is the central question we study in this paper: One wishes to estimate the conditional distribution p(#2 21) from a given flow model p(æ)) for some partitioning of variables (æ), 222). Existing methods for this task largely fall under two categories: Markov Chain Monte Carlo (MCMC) and variational inference (VI). While MCMC methods can perform exact conditional sampling in theory, they often have prohibitively long mixing time for complex high-dimensional distributions and also
For example, Kingma& & Dhariwal 2018 report that their largest model had 200M parameters and was trained 40 GPUs for week.

---

Under review as a conference paper at ICLR 2021
do not provide likelihoods. On the other hand, VI allows for approximate likelihood evaluation under the variational posterior and fast sampling, but at a lower sample quality compared to MCMC counterparts.
We propose novel method that leverages a powerful pre-trained flow model by constructing carefully designed latent codes to generate conditional samples via variational inference. While this procedure intractable for latent variable models in general, the invertibility of the pre-trained model gives us a tractable algorithm for learning distribution in the latent space whose samples approximately match the true conditional when fed into the pre-trained model.
Our contributions:
We start with interesting theoretical hardness result. We show that even though flow models are designed to provide efficient inversion and sampling, even approximate sampling from the exact conditional distribution is provably computationally intractable for wide class of flow models. This motivates our approach of smoothing the observation.
We develop a method to estimate the target conditional distribution by composing second flow model (which we call the pre-generator) with the given model. In particular, this parametrization allows us to employ variational inference and avoid unstable adversarial training as was explored in existing work.
The resulting approximate posterior retains the computational flexibility of a flow model and can be used on a wide variety of downstream tasks that require fast sampling, exact likelihood evaluation, or inversion. Compared to MCMC methods, it has the benefit of being able to generate samples that are guaranteed to be i.i.d.
We experimentally show that our approach is comparable to simple MCMC baselines in terms of sample quality metries such as Frechet Inception Distance (Heusel et al. 2017). We also demonstrate that it achieves superior conditional likelihood estimation performance compared to regular variational inference.
We extend and validate our method for conditioning under arbitrary differentiable transformations with applications to inverse problems. We qualitatively demonstrate its flexibility on various complex inference tasks.
BACKGROUND
2.1 NORMALIZING FLOWS
Normalizing flow models (also known as invertible generative models) represent complex probability distributions by transforming a simple input noise z (typically standard Gaussian) through differen- tiable bijection f : Rdd > Rdd Since f invertible, change of variables formula allows us compute the probability density of x = f(z):
log p(x) = log p(=) + log det df
(a))
dx
where da denotes the Jacobian of the inverse transformation f-1 : x - z. Flow models are explicitly designed that the above expression be easily computed, including the log-determinant term. This tractability allows them to be directly trained with maximum likelihood objective on data.
dr-
Starting from the early works of Dinh al. (2015) and Rezende & Mohamed (2015), there has been extensive research on invertible architectures for generative modeling. Many of them work by composing a series of invertible layers, such in RealNVP (Dinh al. (2016), IAF (Kingma et al.) 2016). Glow (Kingma & Dhariwal 2018), invertible ResNet (Behrmann et al. 2019), and Neural Spline Flows (Durkan et al. 2019).
One of the simplest invertible layer construction is additive coupling layer introduced by Dinh et al. (2015), which served as basis for many other subsequently proposed models mentioned above. In an additive coupling layer, the input variable is partitioned as æ = (æ1, R4 R2 The layer parametrized by a neural network g(x,) : 2 - d used additively transform æ2. Thus the layer's output y (y1, y2) E Red x IR d2 and its inverse can be computed as follows:
= 20
x
y2
= .2 + g(x1))
32 -
2

---

Under review as a conference paper at ICLR 2021
f
f
Figure I: A flow chart of our conditional sampler. First the noise variable is sampled from (0, Ia). This is fed into our pre-generator to output structured noise 2, which is driving the pre-trained base model f to generate conditional samples æ. The central idea is that the pre-generator must produce structured noise that looks Gaussian (so that the samples realistic) but also make the pre-trained hase model f produce samples that satisfy the conditioning. The final conditional sampler is thus defined by the composition of two flow models.
Notably, the determinant of the Jacobian of this transformation is always for any mapping g-
2.2 VARIATIONAL INFERENCE
Variational inference (VI) a family of techniques for estimating the conditional distribution of unobserved variables given observed ones (Jordan et al. 1999 Wainwright et al. 2008 Blei et al.)
2017). At its core, VI tries to find a tractable approximation of the true conditional density by minimizing the KL divergence between them over set of distributions called the variational family.
Within our context of probabilistic inference from a joint distribution p(z), x2)), the KL minimization problem we solve is the following:
arg min DEL (g(22)
p( m1 = x;)),
gé@
where xj is some fixed observation and the variational family Q is chosen such that efficient sampling and likelihood evaluation are possible for all q Q. Note that q is tied to particular of and is not shared across different observations.
While the conditional density p(a² æ = æi) is generally intractable compute, we can efficiently evaluate the joint density p(#1, 222) when it is given as a flow model. Thus, we instead perform variational inference on the conditional joint distribution:
min DKL (4(x2) || p(#2 21 = x;)) = min Blogg g(x2) - log p(æ1 = xi. x2)]
ge
which obtained by dropping the log p(æ) xxj term that is constant w.r.t. g. We note that this procedure is intractable for other types of latent variable models such as VAE or GAN.
HARDNESS OF CONDITIONAL SAMPLING ON FLOWS
Before we present method, first establish hardness result for exact conditional sampling for flow models that use additive coupling layer. Specifically, if an algorithm is able to efficiently sample from the conditional distribution of flow model with additive coupling layers, then this algorithm can be used to solve NP-complete problems efficiently.
Theorem 1. (Informal) Suppose there is an efficient algorithu that can draw samples from the conditional distribution p(æ2 x1) of a normalizing flow model p( implemented with additive coupling layers as defined in Dinh el al. (2015). Then RP = N..
Moreover, even approximating the true conditional distribution remains hard as long as we require that the conditioning is exact. The formal statement of the above theorem and its corollary to the approximate case can found in Appendix A
Importantly, this result implies that allowing approximate sampling does not affect the hardness of the problem, as long as the hard constraint of exact conditioning (æ, = x;) is there. Thus we are motivated to instead consider approximate conditioning where the conditioned variable æ is not required to match the given observation exactly. We also note that the class of flow models that include additíve coupling layers encompasses a large number of existing models (e.g. most of the models in Section 2.1). Thus our hardness result applies to a variety flow models used in practice.
4 APPROXIMATE PROBABILISTIC INFERENCE WITH COMPOSED FLOWS
Notation. Let pf())) denote the base model, a fixed flow model defined by the invertible mapping f + x. The pre-generator (z(2) is similarly defined by the invertible mapping f e + z and represents a distribution in the latent space of the base model. We assume that all flow models use standard Gaussian prior, i.e. z ~ ~N(0,, Ia) - x = f(x) and e (O, Ta) - z = f(e). By composing the base model and the pre-generator, we obtain the composed model p (x) whose samples generated via (O, Ia) f(((e)) Figure details this sampling procedure.
3

---

Under review as a conference paper at ICLR 2021
1
xi)
P(= ã = xi)
p(z p P
Po
1
8
Po
T
2/2
x
T
y
Po
(a) Ambient VI
(b) Latent VI
(c) Latent VI with transformation
Figure 2: Graphical models depicting different ways we perform variational inference. Solid arrows represent the generative direction, and dotted arrows indicate the variational direction.
Our method. Following the principle of leveraging the favorable geometry of the latent space commonly employed in VI and MCMC literature, we also perform our inference in the latent space of the base model. Because the base model maps Gaussian noise samples, the pre-generator can simply focus on putting probability mass the regions of the latent space corresponding to the conditional samples we wish to model. This in contrast with regular variational inference, which essentially trains an entirely new generative model from scratch to approximate the target conditional density.
Motivated by this observation and the hardness of conditional sampling, we propose to perform variational inference the latent space on a smoothed version of the problem with Gaussian smoothing. Specifically, we create smoothed variable क. distributed according to Pa (@: ax) ((; 1 , Ta) and condition on Φ æi instead. Here, a is the smoothing parameter that controls the tightness of this relaxation. We have experimented with non non-Gaussian smoothing kernels such as the one implicitly defined by the perceptual distance measure LPIPS (Zhang et al. 2018). While we chose to use Gaussian smoothing because our preliminary results showed no appreciable difference, we note that a more extensive tuning of the smoothing kernel may lead to a better sample quality, which we leave for future work.
Thus the objective we would like to minimize the KL divergence between our variational distribution and the smoothed conditional density:
DKL (22))(22)) ) py ã = x))).
(1)
For comparison, a direct application of variational inference the image space (which we refer to as Ambient V/) would minimize the following objective:
DKL (Pg || py(22 | =
where we write Py denote the variational distribution that directly models f(x ÿ y²) in the ambient (image) space, parametrized by an invertible mapping z + x.
Intractability of equation [ Note that in our setting, we cannot directly approximate
( T = æij because we only have access to the joint distribution py (y x2) through the base model. Fortunately, the VI loss for approximating the joint conditional p 21 xx) an upper bound to eq. (I)
(PL yor (x) p x ( py = x(),
which we prove Appendix B.2 Thus we are justified our use of the joint VI loss instead of eq. (T). Moreover, this concern vanishes for the general case where we observe transformation of æ, as discussed below. Figures 2a and 2b show the graphical models reflecting this formulation.
Conditioning under differentiable transformations. The flexibility of VI allows us to easily extend our method to conditioning under an arbitrary transformation. Concretely, let T T(x) be a differentiable function and y be some fixed observation in the range of T', We now observe y = T(x) instead, so we similarly define a smoothed variable ý distributed according to Pa (i y) N (g: y, Ia)). We estimate the conditional density py
v²) by minimizing the following objective:
Lours (f) A DKL (o p )
p v²))
1
(p (z)) + Exppy 202
[T((()) - -
where pr denotes the prior distribution the base model. i.e. N/ It). We provide the derivation of eq. (3) in Appendix B. See Figure 2c for comparison to Equation (T)
4

---

Under review as a conference paper at ICLR 2021
This loss function offers an intuitive interpretation. The first term tries to keep the learned distribution Pfof close to the base distribution by pushing Pj to match the prior of the base model, while the second term tries to match the observation y" Moreover, the above expectation can be rewritten in terms of c, which allows to employ the reparametrization trick to obtain low-variance gradient estimator for training:
Cous(() Bxmpy logp)(2) logps(z) + 20t
1 [T(f(z)) - - v2112
(4)
-Bendi0,1a) [ogP)(f(6)) -logp( +
1 [((((()))) -
20
RELATED WORK
Conditional Generative Models. There has been a large amount of work on conditional generative modeling, with varying levels of flexibility for what can be conditioned
In the simplest case, fixed set of observed variables can be directly fed into the model as an auxiliary conditioning input
(Mirza & Osindero 2014 Sohn et al. 2015 Ardizzone al. 2019). Some recent works proposed to extend existing models to support conditioning on arbitrary subsets of variables (Ivanov et al. 2018))
Belghazi et al. 2019 Li al.) 2019). This is much harder task as there are exponentially many subsets of variables that can be conditioned on.
More relevant to our setting is (Engel et al. 2017), which studied conditional sampling from non- invertible latent variable generators such as VAE and GAN. It proposes to adversarially train a GAN-style generator within the latent space of a pre-trained latent variable model, thereby avoiding the issue of intractability of variational inference for non-invertible models. Due to adversarial training and the lack of invertibility of the base model, however, the learned conditional sampler loses the computational flcxibility of our method.
We highlight several reasons why one might prefer our approach over the above methods: (1) The data used to train the given model may not be available, and only the generative model itself is made public. (2) The given model is too costly to train from scratch. (3) We wish to perform custom downstream tasks (e.g. compression, inversion) that are difficult to do with other parametrizations of the posterior (4) We need to condition on a transformation of the variables, instead of just subset of them. (5) We want get some insight on the distribution defined by the given model.
Markoy Chain Monte Carlo Methods. When one is only concerned with generating samples, MCMC techniques offer a promising alternative. Unlike variational inference and conditional models, MCMC methods come with asymptotic guarantees to generate correct samples. Though directly applying MCMC methods on complex high dithesiienii posteriors parametrized by a neural network often comes with many challenges in practice (Papamarkou et al.) (2019), methods based on Langevin Monte Carlo have shown promising results (Neal al.) [01 Welling & Teh [011 [Sogg & Emmon 2019). Moreover, recent works by Parno & Marzouk 2201 8) and Hoffman al. (2019) also leveraged the favorable geometry of the latent space a flow model to improve mixing of MCMC chains. This
idea was later adapted by Nijkamp et al. (2020) in the context of training energy- basedd models.
While our VI-based method lacks asymptotic guarantees, we hope to show that is empirically competitive with simple MCMC techniques. For this purpose, we consider two baselines: Langevin dynamics and PL-MCMC (Cannella et al., 2020). PL-MCMC is particularly relevant as it also tackles the task of conditional sampling from flow model. Unlike method, it constructs Markov chain in the latent space and uses Metropolis-Hastie to sample from the target conditional distribution with asymptotic guarantees.
Inverse Problems with Generative Prior. In a linear inverse problem, a vector x E RRd generates set of measurements y .Ax E RT, where the number of measurements is much smaller than the dimension: m < d. The goal is to reconstruct the vector æ from y". While in general there (potentially infinitely) many possible values of æ that agree with the given measurements, is possible to identify unique solution when there is an additional structural assumption on æ.
Classically, the simplifying structure was that & is sparse, and there has been extensive work in this setting [Tibshirani] 1996; Candes et al. 2006; Donoho et al., 2006 Bickel et al., 2009; Baraniuk] 2007). Recent work has considered alternative structures, such as the vector x coming from generative model. Starting with Bora et (2017), there has been extensive work in this setting as
5

---

Under review as a conference paper at ICLR 2021
well (Grover & Ermon 2019, Mardani al. 2018 Heckel & Hand 2019; Mixon & Villar 2018
Pandit et al. 2019). In particular, Asim et al. (2019) and Shamshad et al. (2019) utilized flow models to solve inverse problems such as compressive sensing, image deblurring, and image inpainting.
is important to note that the above approaches focus on recovering a single point estimate that best matches the measurements. However, there can be many inputs that fit the measurements and thus uncertainty in the reconstruction. Due to this shortcoming, several recent works focused on recovering the distribution of æ conditioned the measurements (Tonolini al. 2019: Zhang & Jin 2019; Adler & Öktem 2018 2019). We note that our approach differs from these, since they are learning-basc methods that require access the training data. On the contrary, our work attempts to perform conditional sampling using a given pre- pre-trained generative model, leveraging all that previous computation to solve a conditional task with reconstruction diversity.
6 QUANTITATIVE EXPERIMENTS
We validate the efficacy of our proposed method terms of both sample quality and likelihood on various inference tasks against three baselines: Ambient VI (as defined by the loss in Equation (2). Langevin Dynamics, and PL-MCMC. We also conduct our experiments on three different datasets (MNIST, CIFAR-IO, and CelebA-HQ) to ensure that our method works across range of settings.
We report four different sample quality metrics: Frechet Inception Distance (FID), Learned Perceptual Image Patch Similarity (LPIPS), and Inception Score (IS) for CIFAR-I0 (Heusel al. 2017; Zhhng ct al. 2018)) Salimans et al. 2016). While not strictly a measure of perceptual similarity, the average mean squared error (MSE) also reported for completeness.
For all our experiments, we use the multiscale Re architecture (inh et al. 2016) for both the base model and the rre-generator. We use Adam optimizer (Kingma Baj 2014) to optimize the weights of the pre-generator using the loss defined in Equation (3). The images used to generate observations were taken from the test set and were not used train the base models. We refer the reader to Appendix C for model hyperparameters and other details of our experiment setup. Variance
Original
Observed
Conditional Samples
Figure 3: Conditional samples generated by method from observing the upper half of CelebA-HQ faces. We see that our approach is able to produce diverse completions with different jaw lines, mouth shapes, and facial expression.
6.1 IMAGE INPAINTING
We perform inpainting tasks using our approach, where we sample missing pixels conditioned on visible ones. We consider three different conditioning schemes: the bottom half (MNIST), the upper half (C)lebA-HQ). and randomly chosen subpixels (CIFAR-10). For MNIST, we use the smoothing parameter value of o 0.1 and for CIFAR-10 10 and CelebA-HQ,
0.05.
In Figure 3 we see that our approach produces natural and diverse samples for the missing part of the image. The empirical pixelwise variance (normalized and averaged over the color channels) also confirms that, while the observation is not perfectly matched, most the high-variance regions are in the unobserved parts
expect.
We also quantitatively evaluate the quality of the generated samples using widely used sample quality metrics, as shown in Table As we see, method outperforms the baseline methods most

---

Under review as a conference paper at ICLR 2021
of the metrics. Note that PL-MCMC results for CIFAR-10 and CelebA-HQ omitted because it was prohibitively slow to run for hundreds of images, as each MCMC chain required over 20,000 proposals. Cannella et al. (2200) also report using 25,000 proposals for their experiments.
Table 1: Sample quality metrics for image inpainting tasks on different datasets. The best value is bolded for each metric. As shown below, method achieves superior sample quality to all baselines.
MNIST
CIFAR 10 (5-bit)
CelebA-HQ (5-bit)
FID
MSE
LPIPS
FID
1St
MSE
LPIPS
FID
MSE
LPIPS
Ours
4.11
21.67
0.074
41.14
7.189
9.71
0.176
33.61
223.06
0.208
Langevin
14.34
36.51
0.135
47.53
6.732
9.31
0.201
30.33
323.47
0.229
Ambient VI
114.59
65.56
0.290
84.78
5.156
16.74
0.296
289.64
1060.66
0.587
PL-MCMC
21.20
59.89
0.190
N/A
N/A
6.2
LIKELIHOOD ESTIMATION
Next, we evaluate our method on the task of conditional likelihood estimation. By varying the size of the pre-generator, we also test the parameter efficiency of our method comparison Ambient VI. Results are shown in Table 2 we see that our method is able to produce reasonable samples using only about 1% of the base model's parameters, confirming the effectiveness of inference latent space.
Table 2: Conditional likelihood estimation performance (measured in bits per dimension) for different pre- vee-generator sizes on the MNIST imputation task. The first row shows the parameter count of the pre ree-generator relative the base model.
Observations
C
# Parameters
Ours
Ambient VI
Conditional Completions (ours)
1.2%
1.73
6.75
0
1
2
3
4 3
6
7
3
3.2%
1.64
3.17
0
1
2
3
3
6
7
5
9
10.6%
1.52
2.71
0
/
2
3
3
6
7
5
9
39.1%
1.47
2.99
0
/
2
3
4
5
6
7
8
9
QUALITATIVE EXPERIMENTS
Here present interesting application of conditioning under a complex transformation T'. Since the only requirement for T is that it must be differentiable. we can set be another neural network. For example, T' a pre-trained binary classifier for specific attribute, we can extract model conditioned on the presence (or the absence) of that attribute from an unconditional base model.
Extracting Class-conditional Models:
We test this idea on the MNIST dataset. We trained 10 binary classifiers on the MNIST dataset, one for each digit of those classifiers, we were able to extract the class-conditional model py (x Label(ar) z k)). See
0,...,9. to predict whether the given image is & not. By setting T be each
Figure 4a for samples generated from the extracted models.
We additionally test the applicability of our method to linear inverse problems. In Figure 5 we show the conditional samples obtained by our method on three different tasks: image colorization, super-resolution (2x), and compressed sensing with 500 random Gaussian measurements (for reference, CIFAR-10 10 images have 3072 dimensions). We notice that the generated samples look natural, even when they do not match the original input perfectly, again showing method's capability to generate semantically meaningful conditional samples and also provide sample diversity.
Inverse Problems:
7

---

Under review as a conference paper at ICLR 2021
Original
Inpainting Samples
0
2
3
7
5
6
7
_ 9
4
0
/
2
3
4
5
7
7
8
4
0
2
3
4
5
6
7
8
9
9
0
2
3
4
5
6
7
8 8
9 9
0
2
3
4
S
6
(a) Each column contains samples from the learned posterior conditioned
different MNIST classes. We emphasize that the base model was trained without class information.
(b) Conditional samples from image inpainting experiments done with Ambient VI.
Samples from class-conditiona models extracted from the unconditional base model (left) and various failure modes of Ambient VI (right).
Figure 4:
Original
Observed
Conditional Samples
Variance
Colorization
SR (2x)
Compressed Sensing
Figure 5: Results various inverse problem tasks using method.
7.1 WIYY AMBIENT VI FAILS
From Table notice that Ambient VI achieves significantly worse sample quality compared to other methods. The low-quality samples from the image inpainting task in Figure/4]furte further confirm that Ambient is unable to produce good conditional samples, even though the observation matched well. This may seem initially surprising but is natural consequence of the VI objective. Recall that our loss function decomposes into two terms: the KL term and the reconstruction term
Lours (f) (p(())(2)) pf())) + Ez-py
[T(((=)) - y
(5)
20
If we alternatively derive the loss for Ambient VI, we arrive at an analogous objective:
Laubiient (g) (P) || p)()))
+ E Eapp,
202 [T(x) y 1'1
While these two loss functions seem like simple reparametrizations of each other via f. they behave very differently during optimization due to the KL term. Notice that for both loss functions, the first term is the reverse KL divergence between the variational distribution and the base distribution. Because reverse KL divergence places no penalty whenever Pj or Py is zero regardless of py. minimizing the reverse KL is known to have a mode-seeking behavior where variational distribution fits single mode of py and ignores the rest of the support of py (Murphy) 2013] Chapter 21.2 21.2.2).
8

---

Under review as a conference paper at ICLR 2021
In contrast, minimizing the forward KL has zero-avoiding behavior and tries to cover all of py's
support.
For method, this is not problem because the prior dis- tribution of the base model (z(2) ) a standard Gaussian and hence unimodal. However, for Ambient VI, Py (a: is the base model itself and is highly multimodal. This can be empirically seen by visualizing the landscape of log (x)) projected onto a random 2D subspace. In Figure 66, we clearly see that ps (æ) has numerous local maxima. For Ambient VI, the variational distribution collapses into one of these modes.
CONCLUSION
We proposed new inference algorithm for distributions parametrized by normalizing flow models. The need for approx- imate inference is motivated by our theoretical hardness result for exact inference, which is surprising given that it applies invertible models. We also presented detailed empirical evaluation of our method with both quantitative and qualitative results on wide range of tasks and datasets.
Contour plot of log py (4)) around a random point in image space.
Figure
There are numerous directions for further research. For example, our method can be possibly extended other latent-variable generators such as VAEs and GANs (Kingma & Welling 2013 Goodfellow et al. 2014). A significant limitation of our work is that we need to re-train the pre- genereeerttor for each observation. It may be possible avoid retraining by amortizing the pre-generaton over all possible observations. Studying the various trade-offs resulting from such scheme would be an interesting result similar in spirit to Cremer et al. (2018). Overall, we believe that the idea a pre-generator creating structured noise is useful and general method for leveraging pre pre-trained generators to solve new generative problems.
REFERENCES
Jonas Adler and Ozan Öktem. Deep bayesian inversion. arXiv preprint arXiv: 001/8//..59100 2018.
Jonas Adler and Ozan Öktem. Deep posterior sampling: Uncertainty quantification for large scale inverse problems. 2019.
Lynton Ardizzone, Carsten Lüth, Jakob Kruse, Carsten Rother, and Ullrich Köthe. Guided image generation with conditional invertible neural networks. arXiv preprint arXiv: 1907.0239 2019.
Muhammad Asim, Ali Ahmed, and Paul Hand. Invertible generative models for inverse problems: mitigating representation error and dataset bias. arXiv preprint arXiv:1905.11677 2019.
Andrei Atanov, Alexandra Volokhova, Arsenii Ashukha, Ivan Sosnovik, and Dmitry Vetrov. Semi- conditional normalizing flows for semi-supervised learning. arXiv preprint arXiv: 2019.
Richard Baraniuk. Compressive sensing. IEEE signal processing magazine, 24(4), 2007.
Jens Behrmann, David Duvenaud, and Jörn-Henrik Jacobsen. Invertible residual networks. Interna- tional Conference on Machine Learning, 2019.
Mohamed Ishmael Belghazi, Maxime Oquab, Yann LeCun, and David Lopez-Paz. Learning about an exponential amount of conditional distributions. arXiv preprint urXiv:/902.06s: 2019.
Peter Bickel, Ya' acov Ritov, Alexandre B Tsybakov, et al. Simultaneous analysis of lasso and dantzig selector. The Annals Statistics, 37(4):1705-177 2009.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: A review for statisticians.
Journal of the American Statistical Association, 1 112(518):55--87, 2017.
Ashish Bora, Ajil Jalal, Eric Price, and Alexandros Dimakis. Compressed sensing using generative models. In International Conference Machine Learning, pp. 537-546. JMLR. org, 2017.

---

Under review as a conference paper at ICLR 2021
Emmanuel J Candes, Justin K. Romberg. and Terence Tao. Stable signal recovery from incomplete and inaccurate measurements. Communications on Pure and Applied Mathematicc:: A Journal Issued by the Courant Institute of Mathematical Sciences, 5908): 2006.
Chris Cannella, Mohammadreza Soltani, and Vahid Tarokh. Projected latent markov chain monte carlo: Conditional sampling of normalizing flows, 2020.
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders. arXiv preprint arXiv: 003558, 2018.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components
estimation. In International Conference on Learning Representations 2015 workshop track, 2015.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using Real NVP. In International Conference on Learning Representations, 2016.
David L. Donoho et al. Compressed sensing. IEEE Transactions on information theory, 52(4): 1289-1306, 2006.
Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Neural Information Processing Systems, 2019.
Jesse Engel, Matthew Hoffman, and Adam Roberts. Latent constraints: Learning to generate conditionally from unconditional generative models. arXiv preprint arXiv: 2017.
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in neural informa- tion processing systems, pp. 2672-2680, 2014.
Aditya Grover and Stefano Ermon. Uncertainty autoencoders: Learning compressed representations via variational information maximization. International Conference on Artificial Intelligence and Statistics, 2019.
Reinhard Heckel and Paul Hand. Deep decoder: Concise image representations from untrained non-convolutional networks. In International Conference on Learning Representations, 2019.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two time-scale update rule converge local nash equilibrium. In Advances in neural information processing systems, pp. 6626-6637, 2017.
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In International Conference on Machine Learning, pp. 2722-2730, 2019.
Matthew Hoffman, Pavel Sountsov, Joshua V Dillon, Ian Langmore, Dustin Tran, and Srinivas Vasudevan. Neutra-lizing bad geometry in hamiltonian monte carlo using neural transport. arXiv preprint arXiv:1903.03700 2019.
Oleg Ivanov, Michael Figurnov, and Dmitry Vetrov. Variational autoencoder with arbitrary condition- ing. arXiv preprint arXiv:1806.02382. 2018.
Michael Jordan, Zoubin Ghahramani, Tommi Jaakkola, and Lawrence Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233, 1999.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arX1v:1412.6980, 2014.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv:1312.6114, 2013.
arXiv preprint
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolutions, In Neural Information Processing Systems, pp. 10215-10224, 2018.
10

---

Under review as a conference paper at ICLR 2021
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, llya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Neural Information Processing Systems, pp- 4743-4751 2016.
Yang Li, Shoaib Akbar, and Junier B Oliva. Flow models for arbitrary conditional likelihoods. arXiv preprint arXiv:1909.063 06399, 2019.
Morteza Mardani, Qingyun Sun, David Donoho, Vardan Papyan, Hatef Monajemi, Shreyas Vasanawala, and John Pauly. Neural proximal gradient descent for compressive imaging. In Neural Information Processing Systems, pp. 9573-9583 2018.
Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXiv preprint arXiv: 1:Xiv 2014.
Dustin G Mixon and Soledad Villar. Sunlayer: Stable denoising with generative networks. arXiv
preprint arXiv: SEL.::v8003001999 2018.
Kevin P. Murphy. Machine learning probabilistic perspective.
MIT Press, Cambridge.
Mass, [u.a.], 2013. ISBN 9780262018029 0262018020. URL hottps://wwie amazon
026201 000008020/10 www-rrrffssrrr 1 028013111000/r1111222ee222443333
Radford M Neal al. Meme using hamiltonian dynamics. Handbook of markov chain monte carlo, 2(11);2, 2011.
Erik Nijkamp, Ruiqi Gao, Pavel Sountsov, Srinivas Vasudevan, Bo Pang, Song-Chun Zhu, and Ying Nian Wu. Learning energy-based baaed model with flow-based backbone by neural transport meme. arXiv preprint arXiv:2006.06897, 2020.
Aaron Oord, Yazhe Li, Igor Babuschkin, Karen Simonyan, Oriol Vinyals, Koray Kavukeuoglu, George Driessche, Edward Lockhart, Luis Cobo, Florian Stimberg, et al. Parallel wavenet: Fast high-fidelity speech synthesis. In International conference on machine learning, pp. 3918-3926. PMLR, 2018.
Parthe Pandit, Mojtaba Sahraee, Sundeep Rangan, and Alyson K Fletcher. Asymptotics of map inference in deep networks. arXiv preprint
20199
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. arXiv preprint arXiv:1912.02762, 2019.
Theodore Papamarkou, Jacob Hinkle, Todd Young, and David Womble. Challenges in bayesian inference via markov chain monte carlo for neural networks. arXiv, pp. arXiv-1910. 2019.
Matthew D Parno and Youssef Marzouk. Transport map accelerated markov chain monte carlo. SIAM/ASA Journal on Uncertainty Quantification, 6(2):645-68) 2018.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Conference on Machine Learning, 1530-1538, 2015.
Tim Salimans, Ian Goodfellow, Wojeiech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. In Advances in neural information processing systems, pp. 2234-2242, 2016.
Fahad Shamshad, Asif Hanif, and Ali Ahmed. Subsampled fourier ptychography via pretrained invertible and untrained network priors. In NeurlPS 2019 Workshop on Solving Inverse Problems with Deep Networks, 2019.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. Neural Information Processing Systems, pp. 3483-3491, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, pp. 11918-11999 1190,, 2019.
11

---

Under review as a conference paper at ICLR 2021
Robert Tibshirani.
Regression shrinkage and selection via the lasso. Journal of the Royal Statistical Society: Series B (Methodological), 58(1):677-285 1996.
Francesco Tonolini, Ashley Lyons, Piergiorgio Caramazza, Daniele Faccio, and Roderick Murray- Smith. Variational inference for computational imaging inverse problems. arXiv preprint 1rXiv:1904.06264, 2019.
Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends@ in Machine Learning. 1(1-2))1-306 2008.
Patrick Nadeem Ward, Ariella Smofsky, and Avishck Joey Bose. Improving exploration in soft-actor- critic with normalizing flows policies. arXiv preprint arXiv:1906.02771, 2019.
Max Welling and Yee W Teh. Bayesian learning via stochastic gradient langevin dynamics. In Proceedings of the 28th international conference on machine learning (ICML-11), pp. 681-688, 2011.
Chen Zhang and Bangti Jin. Probabilistic residual learning for aleatoric uncertainty in image restoration. arXiv preprint (008:000000001111 2019.
Richard Zhang, Phillip Isola, Alexei A Efros, Eli Shechtman, and Oliver Wang. The unreasonable effectiveness of deep features as a perceptual metric. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 586-595, 2018.
12

---

Under review as a conference paper at ICLR 2021
A PROOF OF HARDNESS RESULTS
A.1 PRELIMINARIES
Boolean variable a variable that takes a value (-1,1}. A literal is Boolean variable x; or its negation -I,- A clause is set of literals combined with the OR operator, e.g., V -22 V 23- A conjunctive normal form formula a set of clauses joined by the AND operator, e.g., (I) V -D2 V (a) A (In -33 V s4). A satisfying assignment is assignment the variables such that the Boolean formula is true.
The 3-SAT problem the problem of deciding if conjunctive normal form formula with three literals per clause has satisfying assignment. We will show that conditional sampling from flow models allows us solve the 3-SAT problem.
We ignore the issue of representing samples from the conditional distribution with finite number of bits. However the reduction is still valid if the samples are truncated a constant number of bits.
A.2 DESIGN OF THE ADDITIVE COUPLING NETWORK
Given conjunctive normal form with m clauses, design a ReLU neural network with 3 hidden layers such that the output is 0 if the input is far from satisfying assignment, and the output is about a large number M if the input is close satisfying assignment.
We will define the following scalar function
(-(r (1- (1-5)) ReL.U (1- (1---(1- 1))--)) -1)
(x) =ReLU
[ (x (x- + ReLU (-
ReeUU
This function is the input is 1, 0 if the input x has |x 1| and is linear interpolation on (1 - E, 1 + =). Note that it can be implemented by a hidden layer of a neural network and linear transform, which can be absorbed in the following hidden layer. See Figure for a plot of this function.
1 1
1
Figure 7: Plot of the scalar function used to construct an additive coupling layer that can generate samples of satisfying 3-SAT assignments.
For each variable æe we create transformed variable v; by applying 3; (x,)) (-(=);). Note [1+ [1+ oo). -1 at x, = -1, at I = 1,
that this function is on (-o0, -1 - []j =] U [[ and smooth interpolation on the remaining values in the domain.
Every clause has at most 8 satisfying assignments. For each satisfying assignment we will create a neuron with the following process: (1) get the relevant transformed values ã;, f... .g. (2) multiply each variable by 1/3 if it is equal to the satisfying assignment and -1/3 if it is equal to -1 in the satisfying assignment, (3) sum the scaled variables, (4) apply the 6, function to the
We will then sum all the neurons corresponding to satisfying assignment for clause C, to get the value cy. The final output is the value M x ReLU(E, Cj (m 1)), where M is large scalar.
13

---

Under review as a conference paper at ICLR 2021
We say that an input to the neural network x corresponds to Boolean assignment x' ff-1.11 for every 24 we have [z; - 2 < E. For E < 1/3, if the imput does not correspond to a satisfying assignment of the given formula, then at least one of the values ej is O. The remaining values of cj are at most 1, so the sum in the output is at most (m - 1), thus the sum is at most zero, so the final output is O. However, if the input is satisfying assignment, then every value of Cj 1, the output is M.
A.3 GENERATING SAT SOLUTIONS FROM THE CONDITIONAL DISTRIBUTION
Our flow model will take in Gaussian noise ; Z ~ N(, 1). The values 21,,.. ,... Ta will be passed through the output. The output variable y will be z f'M ( ,..) where sas is the neural network described in the previous section, and M is the parameter in the output to be decided later.
Let A be all the valid satisfying assignments to the given formula. For each assignment a, we will define Xa to be the region Xa = far Rd l|a xl10 I E), where as above is some constant less than 1/3. Let X = UsEA Xa.
Given an element a' E Xa, we can recreate the corresponding satisfying assignment a. Thus if we have an element of XA. we can certify that there is satisfying assignment. We will show that the distribution conditioned on y = M can generate satisfying assignments with high probability.
We have that
ply = MA) p X y = M) ply M, + p'y M,
If we can show that ply = X < plu M, X x), then have that the generated samples with high probability satisfying assignments.
Note that,
ply= = M, p(y = M Xx)P(XA) ply M XA).
Also notice that if x E XA. then fx) = O. Thus ~ .N(,, and Ply M XA)= = O(exp(-M²/2).
Now consider any satisfying assignment &a- Let X be the region X = {x E Rd ||a 1 1 S 3. Note that for every in this region we have fMi (z) 2 M/2. Additionally, we have that P(X)) = @(m) Thus for any I Xx we have p(Y = M x) 2 exp(-M2/8). We can conclude that
/ p(Y M x)p(z) do 2 exp(-M223 Ogd m)). xg
p(y Ma) 2 = M,X))
For M (o m. we have that p(y = M, x)) exponentially smaller than p(y M, XA). This implies that sampling from the distribution conditioned on = M will return a satisfying assignment with high probability.
A.4 HARDNESS OF APPROXIMATE SAMPLING
Definition 2. The complexity class RP is the class of decision problems with efficient random algorithms that (1) output YES with probability 1/2 if the true answer is YES and (2) output NO with probability if the true answer NO. It is widely believed that RP is a strict subset of NP.
simple extension of the above theorem shows that even approximately matching the true conditional distribution in terms of TV distance is computationally hard. The total variation (TV) distance is defined as drv (p. g) = s lp(E) F.) q( F:) < where is an event. The below corollary shows that it is hard to conditionally sample from a distribution that is even slightly bounded away from 1.
Corollary 3. The conditional sampling problem remains hard even if we only require the algorithm to sample from a distribution q such that drv x = x²) ) )) S 1/poly(d) where d is the dimension of the distribution.
We show that the problem is still hard even if we require the algorithm to sample from a distribution q such that (pv(p(z y )"),4) 2 1/poly(d).
Consider the event XA from above. We saw that p(X, I y = M) 2 - e -dd)). We have that drv y M), q) 2 1 exp(- q(X x)).
14

---

Under review as a conference paper at ICLR 2021
Suppose that the distribution q has ((Aa) 2 1/poly(d). Then by sampling a polynomial number of times from q we sample an element of XA. which allows us to find satisfying assignment. would be able to efficiently solve SAT
Thus if we
efficiently create such distribution, and RP = NP. As we are assuming this is false, we must have q( (Xxa) S 1/poly(d), which implies drv (p(: y M),9) 2 1 (/poly(d).
B
MISSING DERIVATIONS
B.I DERIVATION OF EQUATION
Here we present a detailed derivation of Equation 3 Note that this equality is true up to a constant w.t.t f. which is fine as we use this as the optimization objective.
Cour( A (P of (x)
( c y = y")
log logPref(2) - log y²) + logps(ū y²)
= Eawpyos
Ex-pyaf
logp joj - logP/(x)) - log (U = y¹ xz)
Eawpyoj
logg/oo log
+ PRBoppos | log Po yy 1 1y=
(KL(Prof(zz) (x))) + BE
|- Po = y = T(x)))
c DKL(pp(2) 1y + Egup, 20
[[((f(z) - 101
In (A), we drop the log y") term as is constant w.r.t. f. In (B), we use the conditional independence ÿ Il x y.
In (C),
the invariance of KL divergence under invertible transformation to rewrite the KL
divergence terms of 2.
B.2 JOINT VI MARGINAL VI
We also provide a justification for using the joint VI loss as discussed section 4 Specifically, we show that the joint VI loss is upper bound to the intractable marginal VI loss in eq. [ which wish to minimize. For notational brevity, we write ((x)) and g(z) to denote the variational posterior P Poofle (x) and the pre-generator pz(2), respectively. Then, (Joint KL)
= (KL(g(x) ()( 21 = xi))
By(())))) [log ((11, x22) - logp =;]]]
Eg()), [log ((x2) + log ( x2) log (oggr((2)) = xx - logp/(<x क = x x2)]
= By((2)) [log ( logpy = =
Eeg()))) [B [log 7 x2) (0 क =|
DKL (g(x2) || (x2 =x1) + Eg(,,) [KL((( x2) p xffx))))
2 DKL. (g( || (22 =;i)
= (Marginal KI.),
where the last inequality is due to the nonnegativity of KL. Note that equality holds when (KL(((2)
222) || pf(ai [2 æi, #2)) = 0, i.e. when our variational posterior matches the true conditional.
C EXPERIMENT DETAILS
C.1 OUR ALGORITHM
C.2 EFFECTS THE SMOOTHING PARAMETER
The choice of variance for Gaussian smoothing an important hyperparameter, so we provide some empirical analysis of the effects of a. As shown in Figure Sal large values of cause the samples to ignore the observation, whereas small values lead to unnatural samples as the learned distribution becomes degenerate. Visually, achieve essentially negligible variance the observed
15

---

Under review as a conference paper at ICLR 2021
Algorithm Training the pre-generator for given observation under transformation. We assume that f is an invertible neural network with parameters 6.
1: Input: y" observation we are conditioning on, [w)): differentiable transformation of æ.
2: for i 1 ................................ num sttessteesseeee do
3: for 1 1.... m do
> generate m latent codes from Pf(22
4:
Sample e (N Ia) z(6) + f(e())
5:
6:
end for
7: C
(ggp((2(0)
log Prormal ( (z 2 ||(((((0))) - -
8:
f 0 VoL
gradient step
9: end for
Algorithm 2 Sampling from the approximate posterior from the learned pre-generator.
Input: f: base model, f: pre-generator trained on the given observation, n: number of samples.
2: for i = 1... do
3: Sample e(f) (0, La)
f(f(e(0)) 4: xx(1) = (x) x22
(8) (6)
b Feed the noise through composed flow.
5: end for
6: if Partitioned Case then
7: Return {2 ,..., x
Return only the unobserved portion of the samples.
8: else
9: Return (xx11 (
end if
portion past 0.01, but at the slight degradation in the sample quality. In Figure 8b we also notice that the difference between the true observation (x:²) and generated observation (æ1) stops improving past = 1e-4. We tried annealing from large value to small positive target value to see that would help improve the sample quality at very small values a, but noticed no appreciable difference. In practice, recommend using the largest possible that produces observations that are within the (task-specific) acceptable range of the true observation.
C.3 HYPERPARAMETERS: BASE MODEL AND PRE-GENERATOR
See Table and Table for the hyperparameters used to define the network architectures train them. For the color datasets CIFAR-I0 and CelebA-HQ, we used 5-bit pixel quantization following Kingma
Dhariwal (201). Additionally for CelebA-HQ, we used the same train-test split (27,000/ 27,,000/3,000) of
Kingma & Dhariwal| (2018) and resized the images 64 x 64 resolution. Uncurated samples from the base models are included for reference in Figure
Table Hyperparameters used train the base models used in our experiments.
Base Models
MNIST
CIFAR-10
CelebA-HQ
Image resolution
28 x 28
32 X 32
64
64
Num. scales
3
6
Res. blocks per scale
8
12
Res. block channels
32
64
80
Bits per pixel
8
5
5
Batch size
128
64
32
Learning rate
0.001
0.001
0.001
Num. epochs
200
Test set bits-per-dim
1.053
1.725
1.268
16

---

Under review as a conference paper at ICLR 2021
Observed
0.1
1e-3
1e-6
{
1
2
2
2
0.05
MSB observed part
7
2
2
2
3
9
7
7
2
2
0.03
0
2
2
2
2
0.01
-2
-6
l0gn
(a) Each column contains samples learned conditional sampler at different values of a with pixelwise variance computed using 32 samples.
(b) MSE between # and æ at different values a.
Figure 8: Effect of the smoothing parameter on sample quality and tightness of approximate condi- tioning.
Table 4: Hyperparameters used to define and train the pre-generator for each of our experiments.
Base Models
MNIST
CIFAR-10
CelebA-HQ
Image resolution
28 x 28
32 x 32
64 x 64
Num. scales
3
4
Res. blocks per scale
3
4
3
Res. block channels
32
48
48
Batch size
64
32
3
5
4
0
8
8
6
0
7
0
1
5
/
2
2
3
Figure 9: Unconditional samples from the base models used for our experiments. From left: MNIST, 5-bit CIFAR-10, and 5-bit CelebA-HQ models.
C.4 HYPERPARAMETERS: IMAGE INPAINTING
We randomly chose 900/500/300 images from MNIST/CIFAA--IOCCelb.A. masks defined in Section 6.1 and generated samples conditioned on the remaining parts. FID and other sample quality metrics were computed using conditional samples per test image for VI methods (ours and Ambient VI), 8 for Langevin Dynamics, and 6 for PL-MCMC. We note that using more samples for VI methods do not unfairly affect the result of sample quality evaluation, i.c. there was no appreciable difference when using 8 vs. 32 samples to compute FID. We used more samples simply because is much cheaper for VI methods to generate samples compared to MCMC methods.
test sets, applied
For VI Methods (Ours & Ambient VI)
Learning rate: 1c-3 for MNIST; 4 for the others Number of training steps: 2000 for CelebA-HQ: 1000 for the others
For Langevin Dynamics
17

---

Under review as a conference paper at ICLR 2021
Learning rate: 5e-4 for all datasets
Length of chain: 1000 for CIFAR- 10; 2000 for the others
For PL-MCMC
Learning rate: 5e-4)
Length of chain: 2000 for MNIST
%a 1e-3, pp 0.05
C.5 HYPERPARAMETERS: INVERSE PROBLEMS
Colorization
Compressed Sensing
Compressed Sensing
Super-resolution
Learning rate
5e-4
5e-4
5e-4
5e-4
0.05
0.05
0.05
0.05
Dataset
CelebA-HQ
CelebA-HQ
CIFAR-10
CIFAR-10
Batch size
8
32
32
Number of steps
1000
2000
1000
1000
Engel et al. (2017)
Parno & Marzouk
(2018)
Hoffman et al. (209))
Nijkamp al. (222
18


Reviews:
[{"review_id": "sUFP8ex1yOW", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary:\n \nThis paper uses Variational Inference to query pre-trained flow-models. If the flow variable is $x$, then querries are either conditioned on the part of $x=(x_1,x_2)$, or on a differentiable transformation of $x$. Authors first show that it is not trivial to conduct such queries exactly or approximately for a general class of flow-models. The paper then proposes a framework that affords such querries by working in the latent space--empirical evidence favors the proposed approach over contemporary methods.  \n\n### Strength:\n\nThe problem is well-motivated. The authors outline several instances where one would want to work with a pre-trained model and use it to query over new data. Further, I appreciate the use of proof that the problem is hard to motivate the solution. I also commend the author's effort to collect empirical evidence for their method. I especially like the \"Why Ambient VI fails\" explanation and find the contour plot beautiful. \n\n\n### Concerns:\n\nOne primary with paper is the lack of clarity and overload of notation. This is especially true for section 3.\n\nThe equations 1 and 2, use the distribution $p_{f\\circ \\hat{f}}(x_2)$ and  $p_{\\hat{f}}(x_2)$. The preceding section uses $p_{f\\circ \\hat{f}}(x)$ and  $p_{\\hat{f}}(x)$ for $x = (x_1, x_2)$; it is possible that the authors are referring to the marginal for $x_2$. However, it is not immediately clear from the discussion. Further, there is reference to $y$ which is undefined till that point. More so, if we are talking about marginals, then this necessiates the need to evaluate these marginals.  The authors offer no discussion on this. \n\nIn my understanding of the work, $p_f (x_2| \\tilde{x_1} = x_1^*)$ is approximated with $p_f (x| \\tilde{x_1} = x_1^*)$. Thereafter, we can use $p_f (x| \\tilde{x_1} = x_1^*) \\approx p_f (x_2,x_1) p_\\sigma  (x_1^*|x_1)$ to calculate the eq 1. However, this still leaves me uncertain about the marginal distribution $p_{f\\circ \\hat{f}}(x_2)$. I also believe if the above explanation is true, then it is bit of a leap.\n\nHowever, these concerns vanish for the section with differentiable transformations--here, we do not talk about the partitioning of x, so the expressions are straightforward to evaluate. \n\nThe easiest way to convince me would be to offer a clear explanation of the method and to straighten out the notation. It will be great to get an algorithm--like the one in Appendix C--for the partitioning case. \n\n#### Minor concerns and suggestions:\n\nIn section 2, the last paragraph, the formulation is not unique to the authors' framework. It is the fundamental idea of VI to use ELBO over KL divergence; however, the current presentation makes it feel that this is a novel observation made by the authors. \n\nSection 1, third paragraph: \"VI allows for likelihood evaluation and..\"--I believe the term likelihood has been used to refer to $\\log q$--this is confusing as the term is often reserved for $\\log p(x)$ or $\\log p(x|z)$. I will suggest being unambiguous with terminology.\n\nSection 1, fourth paragraph: \"Specifically, we use variational inference to\nlearn a distribution in the latent space ...\"--I find this sentence hard to parse. How is a distribution \"fed\" to the pre-trained model? (this is more of a writing concern than anything else.)\n\nAuthors can consider using less left margin for bullet points under the heading \"our contributions.\"\n\n### Update after the rebuttal\n\nI think that the ideas in this paper are interesting and can inspire new uses. All of us agree that the problem presented here was important, and there is a lot of work to be done in this domain. However, after reading the discussion with other reviewers and their reviews, I believe the manuscript can benefit from another review round. Specifically, the authors can benefit from a thorough revision of the claims in the paper. Further, I would encourage authors to at least investigate how naive amortization approaches fair (irrespective of the result, authors will develop a stronger case for their line of work.) \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "STKnXsIl2l4", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: \nThis paper is concerned with the approximation of conditional densities in trained normalizing flow models. The author use a variational Bayesian approach to estimate the conditional density of a set of variables given another set. As use cases, the authors present results in image inpainting and colorization. The paper also contains a theoretical result showing that exact conditioning is NP hard in additive normalizing flow models.\n\nPros:\n- Relevant problem.\n- Possibly interesting theoretical result.\n- Interesting range of experiments with careful analysis of both quantitative and qualitative results. Large and appropriate list of performance metrics\n\nCons:\n- Low originality, the main idea to perform inference in the latent space is obvious\n- The relevance of the theoretical result in justifying the rest of the paper is not clear and not enough space is given to its explanation.\n- The smoothing of the likelihood using a very rough noise model is not required\n- More baselines are needed. In particular, it would be useful to have baselines without likelihood smoothing and also comparison with (simulation based) forward KL methods such as [1,3].\n\nRelevance:\nThe problem of conditioning in generative models is highly relevant in our current ML environment as it allows to convert generators into very flexible inference machines capable of solving a large variety of problems. \n\nOriginality: With the possible exception of the NP-hardness theorem, the original contributions of this work are very limited with some questionable elements (see below).  Conditioning deep differentiable models is a standard domain of application of variational Bayesian inference and the authors use a very standard and natural  approach. The main idea behind the method is to perform inference in the latent space, this in my opinion is not a noteworthy contribution as it is just the obvious way to do inference in this setting. The second methodological trick is to use a Gaussian smoothing of the likelihood. To my understanding, this procedure is neither well motivated nor supported experimentally. \n\nMajor concerns:\nI could be missing something but I do not understand why the authors are smoothing the likelihood since the flow already give a perfectly well-behaved joint model. The authors also seem to motivate this choice using their theorem showing that exact conditioning is NP hard. However this is only a valid motivation for using approximate inference, not for adopting an approximate likelihood. It seems that the authors want to get a least-square loss component in the pixel (or add-hoc feature) space. Without smoothing you will get a least square loss in the latent space which is likely to be much more appropriate. Therefore, I am not convinced that this work better in practice and it should at least be tested experimentally. \n\nPaper structure:\nIn general, the paper is well structured with a clear narrative and an appropriate amount of background material. \nHowever, the treatment of the theorem in section 3 should be expanded. The experiment section is very well structured and the analysis of the results is very good and definitely above average. I am pleased that the authors took the time to analyze the bad performance of the native VI method.\n\nWriting:\nThe paper is very clearly written. \n\nLiterature\nThe coverage of the literature is appropriate. However, the author should also discuss methods based on synthetic sampling and forward KL divergence as they are a very viable approach to generative model conditioning in the latent space. For example:\n\n[1] Papamakarios, George, and Iain Murray. \"Fast ε-free inference of simulation models with bayesian conditional density estimation.\" Advances in Neural Information Processing Systems. 2016.\n[2] Le, Tuan Anh, Atilim Gunes Baydin, and Frank Wood. \"Inference compilation and universal probabilistic programming.\" Artificial Intelligence and Statistics. PMLR, 2017.\n[3] Ambrogioni, Luca, et al. \"Forward amortized inference for likelihood-free variational marginalization.\" The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ps8inpQVSYq", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The paper proposes to solve the conditional inference problem by performing a relaxed version of variational inference in the prior space of the flow-based model. The model p(x) is pretrained, and one is interested sampling from p(x|observation). The observation could be some subset of x (inpainting), gray scale representation of an image (colorization), lower-res representation (super resolution), or noisy version of the data (compressed sensing). The paper proposes to perform inference in the prior space, by composing the post-hoc trained latent flow with the trained invertible decoder, as the conditional distribution in that space is believed to have a better geometry.  \n\nContributions and novelties: (A) propose to perform inference in the latent space of a latent variable model to side step the bad geometry, (B) propose to replace hard constraint with stochastic relaxation (placing an additional likelihood term to model the dummy variable). (C) Applications seem interesting for testing the quality of an unconditional flow-based model.\n\nFlaws: missing several important references in the discussion to prior works. These include [1], which describes a more general framework to post-hoc perform sampling from a conditional distribution of a learned latent variable model by fitting a distribution in the latent space; [2,3,4] which propose to mitigate the bad geometry of the learned data energy (in this case, the density model itself) by transforming it into a space where it’s more Gaussianized. Similar idea has also been incorporated in [5] to enable the training of a residual EBM (no need to cite). The key novelties (A and B) are incremental in nature given the above related works that are not cited.\n\n\n[1] Latent Constraints: Learning to Generate Conditionally from Unconditional Generative Models, 2017\n[2] Transport map accelerated markov chain Monte Carlo, 2014\n[3] NeuTra-lizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport, 2019\n[4] Learning Energy-based Model with Flow-based Backbone by Neural Transport MCMC, 2020\n[5] VAEBM: A Symbiosis between Variational Autoencoders and Energy-based Models, 2020\n\nAdditional details:\n\nFor experiments, please report the bits-per-dim (BPD) of the pre-trained flow model, as well as unconditional samples for reference and better comparison. \n“Ambient” langevin dynamics is a very weak baseline, since one can easily improve the mixing via deriving a kernel in the latent space (as per [3]). It wouldl be a more fair comparison since the proposed method also composes with the learned decoder flow. Same goes to the naive implementation of the PL-MCMC which only uses a random walk kernel, which is a very weak baseline.  \nAre the rows of Fig 4 independent? Why do they have the same stroke style if they are all conditioned only on the label only? Is it an indicator of the mode seeking behavior of the reverse KL?   \nIn the qualitative results (tab 2, fig 3 and fig 5), the generated samples do not really satisfy the hard constraints that they are conditioned on (e.g. the subsets of x for inpainting), possibly due to the relaxation via the dummy variable. This is not desirable. Is it possible to anneal the $\\sigma$ while training the latent flow so that it will concentrate on the (potentially degenerate) solution that satisfies these constraints? \nThe authors claim the hardness result is surprising. It has been long shown that sampling from a general Bayesian Belief Network is NP-hard [6]. Can’t the same conclusion be derived from that, with the main difference being the explicit parameterization via a coupling flow? \n\n[6] The Computational Complexity of Probabilistic Inference Using Bayesian Belief Networks, Cooper, 1990\n\n\n\nAdditional questions about the hardness result (I didn’t read the proof):\n\nGenerality of the hardness result: the statement is about the hardness of sampling from the conditional distribution. What is the conditional distribution in this context, is it referring to p(x2|x1) (i.e. observation is a subset of x)? Please be precise. If this is the case it is consistent with the presentation of the previous section (VI). However the proposed method seems to require a more general treatment to take account of the other tasks, e.g. inverse problems. \n\nThe discussion of hardness seems to be used to motivate the relaxation of the hard constraints (of the givens). It doesn’t seem that relevant when the observation is not a deterministic function of x (e.g. inpainting, coloration, etc). For compressed sensing for example, the likelihood p(observation|x) naturally exists and is non-degenerate (by assumption). The presentation seems a bit confusing. \n\n\n--- POST REBUTTAL ---\n\nModified my score after the rebuttal, since (1) I believe the re-purposing achieved by this work can potentially broaden the applicability of flow-based generative model (2) the authors have toned down the abstract and clarified the contributions in the intro, which now better reflects the value of the work. \n\nI am still leaning towards rejection at the end given the limited originality of the proposed method and the lack of a more comprehensive discussion of different possible approaches, but as means to the same end. For example, the relaxed inference problem can be solved with an MCMC method. These should all be discussed and compared if the contribution is about repurposing a joint likelihood model using flows.\n \n\nPS. the last line (the references) of the last page might have been a mistake. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "OUOrPyW2k3j", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "In this work, the authors propose a novel method of estimating conditional distributions over arbitrary partitions of variables $x = [x_1,x_2]$ using an existing pre-trained flow model for $p(x)$. Their method fits a new *pre-generator* flow $\\hat{f}$ for each observation which maps from a base distribution $\\epsilon \\sim N(0,\\mathbb{I})$ to the latent variables $z \\sim N(0,\\mathbb{I})$ where the mapping $z \\leftrightarrow x$ is learned by the pre-trained \"base\" flow $f$. The result is a mapping $\\hat{f}$ which learns to shift probability mass to regions of the latent space which correspond to the conditional distribution of the given observation. The authors present comprehensive experimental results which show a clear improvement over existing methods for conditional inference but with the drawback of needing to re-train the pre-generator for each individual observation.\n\nPros:\n- Very well written, clear, easy to understand\n- The proposed method is intuitive and well defined\n- Placement of the method relative to recent work is very clearly explained\n- Comprehensive theoretical analysis including an interesting hardness result, which is uncommon for the deep learning literature\n- Comprehensive and convincing empirical analysis with clear results\n\nCons:\n- There is very little discussion on how the construction of the base generator $f$ affects the results of the proposed method\n- The proof of hardness is somewhat opaque and feels contrived; but this is often the case with hardness proofs!\n- The method has a clear weakness in needing to be retrained for each observation. However, this is clearly stated by the authors and left open as a direction for future work.\n\nOverall, I think this is an exceptional paper which makes a significant contribution to the field. I think it is suitable to accept as-is with only a few minor adjustments which I will enumerate below.\n\n1. It would be nice (but not absolutely necessary) to see some discussion regarding the construction of the base generator, as I mentioned in the Cons above; e.g. does the performance of this method depend significantly on the user's choice of base model? Intuitively, I would think so.\n2. A few notes on the proofs:\n- The variable lower-case $m$ shows up in several spots in the hardness proof but is never defined. Perhaps these are typos and you meant to write $M$?\n- In the proof for equation 3, the notation for expectations (i.e. $\\mathbb{E}$) is inconsistent in a few places. Presumably just typos.\n- I may be missing something, but it's not immediately clear why $y=T(x)$ can be substituted for the conditioner $x$ in $p_{\\sigma}(\\tilde{y}=y^*|x)$. My immediate intuition is that this would only be valid if $T$ is injective, otherwise this may change the underlying conditional density. Please correct me if I am wrong, and preferably add a clarification to the proof as to why this is justified.\n\nCongratulations to the authors on a job well done!", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a method for conditional inference with arbitrary conditioning by creating composed flows. The paper provides a hardness result for arbitrary conditional queries. Motivated by the fact that conditional inference is hard the paper therefore suggests a novel relaxation where the *conditioning* is relaxed.

There were various concerns from the reviewers regarding notation, comparison algorithms, and how the hardness result motivates the smoothing operation introduced. After careful study of the paper and all the comments I find that I am most concerned about the hardness result and how it motivates the smoothing operation that is done. Novel computational complexity results *as such* are not really in the scope of ICLR. There's nothing wrong with having such a result in a paper, of course, but a paper like this should be evaluated on the basis of the algorithm proposed.

Like R4, I do not follow how this hardness result is meant to motivate the smoothing that's applied. The paper is unambiguous that the goal is to do conditional inference. A hardness result is presented for conditional inference, and so a relaxed surrogate is presented. This has a minor problem that it's not clear the relaxed problem avoids the complexity boundary of the original one. There's a larger problem, though. The hardness result has not been sidestepped! The goal is still to solve conditional inference. The algorithm that's presented is still an approximate algorithm for conditional inference. R4 suggests that other approximation algorithms should be compared to. The authors responded to this point, but I am not able to understand the response. For the same reason, I think it is valid to ask for comparison to other approximate inference algorithms (e.g. without smoothing)

None of the above is to say that the smoothing approach is bad. It may very well be. However, I think that either the existing argument should be clarified or a different argument should be given.

Finally here are two minor points (These weren't raised by reviewers and aren't significant for acceptance of the paper. I'm just bringing them up in case they are useful.)

Is Eq. 3 (proof in Appendix B.1) not just an example of the invariance of the KL-divergence under diffeomorphisms?

Proof in appendix B.2 appears to just a special case of the standard chain rule of KL-divergence (e.g. as covered in Cover and Thomas)
Decision: Reject