Approximate Probabilistic Inference with Composed Flows
Abstract
We study the problem of probabilistic inference on the joint distribution deﬁned by a normalizing ﬂow model.
Given a pre-trained ﬂow model p ( x ) , we wish to estimate p ( x 2 | x 1 ) for some partitioning of the variables x = ( x 1 , x 2 ) .
We ﬁrst show that this task is computationally hard for a large class of ﬂow models.
Motivated by this, we propose a framework for approximate probabilistic inference.
Speciﬁcally, our method trains a new ﬂow model with the property that its composition with the given model approximates the target conditional distribution.
We describe how we can train this new model using variational inference and handle conditioning under arbitrary differentiable transformations.
Experimentally, our approach outperforms Langevin Dynamics in terms of sample quality, while requiring much fewer parameters and training time compared to regular variational inference.
We further validate the ﬂexibility of our method on a variety of inference tasks with applications to inverse problems.
Introduction
Deep generative models have seen an unprecedented growth in the recent years.
Among them, normalizing ﬂow models [1] stand out due to their computational ﬂexibility, as they offer efﬁcient sampling, likelihood evaluation and inversion.
Even with such computational ﬂexibility, however, efﬁcient probabilistic inference on a ﬂow model still remains largely unsolved.
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 practitioners 1 .
If it was possible to perform probabilistic inference on ﬂow models, we could re-purpose these powerful pre-trained generators for numerous custom tasks.
We propose a novel method that leverages a powerful pre-trained ﬂow model by constructing carefully designed latent codes to generate conditional samples.
Speciﬁcally, we use variational inference to learn a distribution in the latent space of the given model such that, when fed into the pre-trained model, the output approximately matches the samples from the true conditional.
Our contributions:
• We establish a hardness result that, even though ﬂow models are designed to provide efﬁcient inversion and sampling, exact conditional sampling is provably computationally intractable for a wide class of ﬂow models.
This motivates our approach of approximate inference.
• We develop a method to estimate the target conditional distribution by composing a second generative model (which we call the pre-generator ) with the given model.
Speciﬁcally, we show how to train the pre-generator to yield structured noise so that the composed model approximately matches the target conditional distribution.
• We experimentally show that our approach is competitive with MCMC methods in terms of speed and widely-used sample quality metrics.
We also demonstrate that it achieves superior conditional likelihood estimation performance compared to regular variational inference.
We further extend and validate the ﬂexibility of our method on conditioning under arbitrary differentiable transformations with applications to inverse problems.
Hardness of Conditional Sampling on Flows
Before we present our method, we ﬁrst establish a hardness result for exact conditional sampling for ﬂow models that use additive coupling layer.
Speciﬁcally, if an algorithm is able to efﬁciently sample from the conditional distribution of a ﬂow model with additive coupling layers, then this algorithm can be used to solve NP-complete problems efﬁciently.
Our hardness result holds even if we allow the conditional sampler to approximately match the conditional distribution.
The formal statement of the below theorem and its corollary to the approximate case can be found in Appendix A.
Theorem 1.
(Informal)
Suppose there is an efﬁcient algorithm that can draw samples from the conditional distribution of a normalizing ﬂow model implemented with additive coupling layers as deﬁned in Dinh et al.
Then RP = NP .
Approximate Probabilistic Inference with Composed Flows
Notation.
Let p f ( x ) denote the base model , a ﬁxed ﬂow model deﬁned by the invertible mapping f : z (cid:55)→ x .
The pre-generator p ˆ f ( z ) is similarly deﬁned by the invertible mapping ˆ f : (cid:15) (cid:55)→ z and represents a distribution in the latent space of the base model.
We assume that all ﬂow models use standard Gaussian prior, i.e. z ∼ N ( 0 , I d ) → x = f ( z ) and (cid:15) ∼ N ( 0 , I d ) →
z = ˆ f ( (cid:15) ) .
By composing the base model and the pre-generator, we obtain the composed model p f ◦ ˆ f ( x ) whose samples are generated via (cid:15) ∼ N ( 0 , I d ) → x = f ( ˆ f ( (cid:15) )) .
Our method.
Motivated by the hardness of exact conditional sampling, we propose to perform variational inference in the latent space on a smoothed version of the problem where we allow the observed variable to approximately match the given observation.
We do this by creating a dummy variable ˜ x 1 distributed according to p σ (˜ x 1 | x 1 ) =
N (˜ x 1 ; x 1 , σ 2 I d ) .
Here, σ is the smoothing parameter that controls the tightness of this relaxation.
The objective we minimize is the KL divergence between our variational distribution and the smoothed conditional density:
[Equation]
For comparison, a direct application of variational inference in the image space (which we refer to as
Ambient VI ) would minimize the following objective:
[Equation]
where we write p ˆ f to denote the variational distribution which directly models f ( x | ˜
y = y ∗ ) .
Note that in our setting, we cannot directly model p f ( x 2 | ˜ x 1 = x ∗ 1 ) because we only have access to the joint distribution p f ( x 1 , x 2 ) through the base model.
Thus we instead approximate the joint distribution conditioned on our observation, i.e. p f ( x | ˜ x 1 = x ∗ 1 ) .
Figures 1a and 1b show the graphical models describing this formulation.
Conditioning under differentiable transformations.
The ﬂexibility of VI allows us to easily extend our method to conditioning under an arbitrary transformation.
Concretely, let T ( x ) be a differentiable function and y ∗ be some ﬁxed observation in the range of T .
We now observe y = T ( x ) instead, so we similarly deﬁne a dummy variable ˜
y distributed according to p σ (˜ y | y )
=
N (˜ y ; y , σ 2 I d ) .
We estimate the conditional density p f ( x | ˜ y
= y ∗ ) by minimizing the following objective:
[Equation]
z xx˜ x 1 f p σ p ( x | ˜ x 1 = x ∗ 1 ) (a
[Figure]
Ambient VI z xx˜ x 1 f p σ p ( z | ˜ x 1 = x ∗ 1 ) (b) Latent VI z x y ˜ y f
T p σ p ( z | ˜ y = y ∗ ) [Figure] [Figure]
We see that our approach is able to produce diverse completions with different jaw lines, mouth shapes,
[Figure]
where p f ( z ) denotes the prior distribution of the base model, i.e. N ( 0 , I d ) .
We provide the derivation of eq.
(3) in Appendix C. See Figure 1c for a comparison to eq.
Notice that the above expectation can be rewritten in terms of (cid:15) to employ the reparametrization trick to obtain a low-variance gradient estimator for training.
Quantitative Experiments
We validate the efﬁcacy of our proposed method in terms of both sample quality and likelihood on various inference tasks against three baselines: Ambient VI (as deﬁned by the loss in eq.
Langevin Dynamics, and PL-MCMC.
We also conduct our experiments on three different datasets (MNIST, CIFAR-10, and CelebA-HQ) to ensure that our method works across a 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-10 [4–6].
While not strictly a measure of perceptual similarity, average mean squared error (MSE) is also reported for completeness.
For all our experiments, we use the multiscale RealNVP architecture [7] for both the base model and the pre-generator.
We use Adam optimizer [8] to optimize the weights of the pre-generator using the loss deﬁned in Equation (3).
The images used to generate observations were taken from the test set and were not used to train the base models.
We refer the reader to Appendix D for model hyperparameters and other details of our experiment setup.
[Table]
4.1 Image Inpainting
We perform inpainting tasks using our approach, where we sample missing pixels conditioned on the visible ones.
We consider three different conditioning schemes: the bottom half (MNIST), the upper half (CelebA-HQ), and randomly chosen subpixels (CIFAR-10).
For MNIST, we use the smoothing parameter value of σ = 0 .
1 and for CIFAR-10 and CelebA-HQ, we use σ = 0 .
In Figure 2a 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 conﬁrms that, while the observation is not perfectly matched, most of the high-variance regions are in the unobserved parts as we expect.
We also quantitatively evaluate the quality of the generated samples using widely used sample quality metrics, as shown in Table 1.
As we can see, our method outperforms the baseline methods on most of the metrics.
Note that PL-MCMC results for CIFAR-10 and CelebA-HQ are omitted because it was prohibitively slow to run for hundreds of images, as each MCMC chain required over 20, proposals.
Cannella et al.
[9] also report using 25,000 proposals for their experiments.
Observing Table 1 closely, one might notice that Ambient VI performs signiﬁcantly worse than other methods.
While it may seem strange at ﬁrst, this is actually a natural consequence of the Ambient VI formulation (see eq.
We provide a detailed explanation for this in Appendix B.
4.2 Likelihood Estimation
Next, we evaluate our method on the task of conditional likelihood estimation.
By varying the size of the pre-generator, we observe the parameter efﬁciency of our method compared to Ambient VI.
Results are shown in Table 2; we see that our method produces reasonable samples using only about 1% of the base model’s parameters, conﬁrming the effectiveness of inference in the latent space.
[Table]
Qualitative Experiments
Extracting Class-conditional Models:
Here we present an interesting application of conditioning under a transformation T parameterized by a neural network.
If T is a pre-trained binary classiﬁer for a speciﬁc attribute, we can extract a model conditioned on the presence (or the absence) of that attribute from an unconditional base model.
We test this idea on the MNIST dataset.
We train 10 binary classiﬁers, one for each digit k = 0 , . . .
, 9 , to predict whether the given image is k or not.
By setting T to be each of those classiﬁers, we are able to extract the class-conditional model p f ( x | Label ( x ) ≈ k ) .
See Figure 3 for samples generated from the extracted models.
Inverse Problems: We additionally test the applicability of our method to linear inverse problems on CIFAR-10 and CelebA-HQ images.
In Figure 2b, we show the conditional samples obtained by our method on three different tasks: image colorization, super-resolution ( 2 × ), and compressed sensing with 500 random Gaussian measurements.
We notice that the generated samples look natural, even when they do not match the original input perfectly, again showing our method’s capability to generate semantically meaningful conditional samples and also provide sample diversity.
Conclusion
We proposed a new inference algorithm for distributions parametrized by a ﬂow.
The need for approximate inference is motivated by the hardness of exact inference.
We also presented a detailed empirical evaluation of our method with both quantitative and qualitative results on a wide range of tasks and datasets.
Overall, we believe that the idea of a pre-generator creating structured noise is a useful and general method for leveraging pre-trained generators to solve new generative problems.
Acknowledgments
This research has been supported by NSF Grants CCF 1763702,1934932, AF 1901292, 2008710, 2019844 research gifts by Western Digital, WNCG IAP, computing resources from TACC and the
Archie Straiton Fellowship.
A Proof of Hardness Results A.1 Preliminaries
A Boolean variable is a variable that takes a value in {− 1 , 1 } .
A literal is a Boolean variable x i or its negation ¬ x i .
A clause is set of literals combined with the OR operator, e.g., x 1 ∨ ¬ x 2 ∨ x 3 .
A conjunctive normal form formula is a set of clauses joined by the AND operator, e.g., ( x 1 ∨ ¬ x 2 ∨ x 3 ) ∧ ( x 1 ∨ ¬ x 3 ∨ x 4 ) .
A satisfying assignment is an assignment to the variables such that the
Boolean formula is true.
The 3-SAT problem is the problem of deciding if a conjunctive normal form formula with three literals per clause has a satisfying assignment.
We will show that conditional sampling from ﬂow models allows us to solve the 3-SAT problem.
We ignore the issue of representing samples from the conditional distribution with a ﬁnite number of bits.
However the reduction is still valid if the samples are truncated to a constant number of bits.
A.2 Design of the Additive Coupling Network
Given a Boolean formula, we design a ReLU neural network with 3 hidden layers such that the [-outputis 
-]output is 0 if the input is far from a satisfying assignment, and the output is about a large number M if the input is close to a satisfying assignment.
We will deﬁne the following scalar function
[Equation]
This function is 1 if the input is 1 , 0 if the input x has | x − 1 | ≥ ε and is a linear interpolation on (1 − ε, 1 + ε ) .
Note that it can be implemented by a hidden layer of a neural network and a linear transform, which can be absorbed in the following hidden layer.
See Figure 4 for a plot of this function.
For each variable x i , we create a transformed variable ˜
x
i by applying ˜ x i =
δ ε ( x i ) − δ ε ( − x i ) .
Note that this function is 0 on ( −∞ , − 1 − ε ] ∪ [ − 1 + ε, 1 − ε ] ∪ [1 + ε, ∞ ) , − 1 at x i = − 1 , 1 at x i = 1 , and a 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
˜ x i , ˜ x j , ˜ x k , (2) multiply each variable by 1 / 3 if it is equal to 1 in the satisfying assignment and − 1 / 3 if it is equal to − 1 in the satisfying assignment, (3) sum the scaled variables, (4) apply the δ ε function to the sum.
We will then sum all the neurons corresponding to a satisfying assignment for clause C j to get the value c j .
The ﬁnal output is the value M × ReLU( (cid:80)
j c j − ( m − 1)) , where M is a large scalar.
We say that an input to the neural network x corresponds to a Boolean assignment x (cid:48) ∈ {− 1 , 1 } d if for every x i we have | x i − x (cid:48)
i | < ε .
For ε < 1 / 3 , if the input does not correspond to a satisfying assignment of the given formula, then at least one of the values
c j is 0 .
The remaining values of c j are at most 1 , so the sum in the output is at most ( m − 1) , thus the sum is at most zero, so the ﬁnal output is 0 .
However, if the input is a satisfying assignment, then every value of c j = 1 , so the output is M .
A.3 Generating SAT Solutions from the Conditional Distribution
Our ﬂow model will take in Gaussian noise x 1 , . . .
, x d , z ∼ N (0 , 1) .
The values x 1 , . .
, x d will be passed through to the output.
The output variable y will be z + f M ( x 1 , . . .
, x d ) , where f M 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 deﬁne X a to be the region X a = { x ∈ R
d : (cid:107) a − x (cid:107) ∞ ≤ ε } , where as above ε is some constant less than 1 / 3 .
Let X A
=
(cid:83) a ∈
A X a .
Given an element x ∈ X a , we can recreate the corresponding satisfying assignment a .
Thus if we have an element of X A , we can certify that there is a satisfying assignment.
We will show that the distribution conditioned on y =
M can generate satisfying assignments with high probability.
We have that
[Equation]
If we can show that p ( y = M, X A ) (cid:28) p ( y = M, X A ) , then we have that the generated samples are with high probability satisfying assignments.
Note that,
[Equation]
Also notice that if x ∈ X A , then f M ( x ) = 0 .
Thus y ∼ N (0 , 1) and P ( y =
M | X A ) =
Θ(exp( − M 2 / 2)) .
Now consider any satisfying assignment x a .
Let X (cid:48) a be the region X (cid:48) a = { x ∈ R d : (cid:107) a − x (cid:107) ∞ ≤ 12 m } .
Note that for every x in this region we have f M ( x ) ≥ M/ 2 .
Additionally, we have that
P ( X (cid:48) a ) =
Θ( m ) − d .
Thus for any x ∈ X (cid:48) a , we have p ( Y = M | x ) (cid:38) exp( − M 2 / 8) .
We can conclude that
[Equation]
For M = O ( √ d log m ) , we have that p ( y = M, X A ) is exponentially smaller than p ( y = M, X A ) .
This implies that sampling from the distribution conditioned on y =
M will return a satisfying assignment with high probability.
A.4 Hardness of Approximate Sampling
Deﬁnition 2.
The complexity class RP is the class of decision problems with efﬁcient random algorithms that (1) output YES with probability 1 / 2 if the true answer is YES and (2) output NO with probability 1 if the true answer is NO.
It is widely believed that RP is a strict subset of NP .
A 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 deﬁned as d TV ( p, q ) = sup E | p ( E ) − q ( E ) | ≤ 1 , where E 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 d TV ( p ( · | x = x ∗ ) , q ) ≤ 1 − 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 d TV ( p ( x | y =
y ∗ ) , q ) ≥ 1 / poly( d ) .
Consider the event
X
A from above.
We saw that p ( X A | y
=
M ) ≥ 1 − exp( − Ω( d )) .
We have that d TV ( p ( · | y = M ) , q ) ≥ 1 − exp( − Ω( d ) − q ( X A )) .
Suppose that the distribution q has q ( X A ) ≥ 1 / poly( d ) .
Then by sampling a polynomial number of times from q we sample an element of X A , which allows us to ﬁnd a satisfying assignment.
Thus if we can efﬁciently create such a distribution, we would be able to efﬁciently solve SAT and RP = NP.
As we are assuming this is false, we must have q ( X A ) ≤ 1 / poly( d ) , which implies d TV ( p ( · | y = M ) , q ) ≥ 1 − 1 / poly( d ) .
B Why Ambient VI Fails
From Table 1, notice that Ambient VI achieves signiﬁcantly worse sample quality compared to other methods.
The low-quality samples from the image inpainting task in Figure 5b further conﬁrm that
Ambient VI is unable to produce good conditional samples, even though the observation is matched well.
This may seem initially surprising but is a natural consequence of the VI objective.
Recall that our loss function decomposes into two terms: the KL term and the reconstruction term.
[Equation]
If we alternatively derive the loss for Ambient VI, we arrive at an analogous objective:
[Equation]
[Figure]
Original
[Table]
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 ﬁrst term is the reverse KL divergence between the variational distribution and the base distribution
D KL ( p ˆ f (cid:107) p f ) .
Because reverse KL divergence places no penalty whenever p ˆ f is zero regardless of p f , minimizing the reverse KL is known to have a mode-seeking behavior where p ˆ f ﬁts a single mode of p f and ignores the rest of the support of p f [10, Chapter 21.2.2].
In contrast, minimizing the forward KL has a zero-avoiding behavior and tries to cover all of p f ’s support.
For our method, this is not a problem because the prior distribution of the base model p f ( z ) is a standard Gaussian and hence unimodal.
However, for Ambient VI, p f ( x ) is the base model itself and is highly multimodal.
This can be empirically seen by visualizing the landscape of log p f ( x ) projected onto a random 2D subspace.
In Figure 5a, we clearly see that p f ( x ) has numerous local maxima.
For Ambient VI, the variational distribution collapses into one of these modes.
C Derivation of Equation (3)
Here we present a detailed derivation of Equation (3).
Note that this equality is true up to a constant w.r.t.
ˆ f , which is ﬁne as we use this as the optimization objective.
[Equation]
In ( A ), we drop the log p f
(˜ y
=
y ∗ ) term as it is constant w.r.t.
ˆ f .
In ( B ), we use the invariance of KL divergence under invertible transformation to rewrite the KL divergence in terms of z .
D Experiment Details D.1 Our Algorithm
1: Input : y ∗ : observation on, T ( x ) : differentiable transformation of x .
2: for i = 1 . . .
num _ steps do 3: for j = 1 . . .
m do (cid:46) generate m latent codes from p ˆ f
( z ) 4:
Sample (cid:15) ( j ) ∼ N ( 0 , I d ) 5:
z ( j ) ← ˆ f ( (cid:15) ( j ) ) 6: end for 7:
L ← 1 m m (cid:80) j = (cid:104) log p ˆ f ( z ( j ) ) − log p normal [Equation] ( j )
[Equation]
T ( f ( z ( j ) )
[Equation]
− y
[Equation] (cid:13)(cid:13) [Equation] (cid:105) 8: θ ← θ − ∇ θ L (cid:46) gradient step 9: end for
D.2 Hyperparameters: Base Model and Pre-generator
See Table 3 and Table 4 for the hyperparameters used to deﬁne the network architectures train them.
For the color datasets CIFAR-10 and CelebA-HQ, we used 5-bit pixel quantization following Kingma and Dhariwal [2].
Additionally for CelebA-HQ, we used the same train-test split (27,000/3,000) of
Kingma and Dhariwal [2] and resized the images to 64 × 64 resolution.
[Table] [Table]
D.3 Hyperparameters: Image Inpainting
We randomly chose 900/500/300 images from MNIST/CIFAR-10/CelebA-HQ test sets, applied masks deﬁned in Section 4.1, and generated samples conditioned on the remaining parts.
FID and other sample quality metrics were computed using 32 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.e. there was no appreciable difference when using 8 vs. 32 samples to compute FID.
We used more samples simply because it is much cheaper for VI methods to generate samples compared to MCMC methods.
For VI Methods (Ours & Ambient VI)
• Learning rate: 1e − 3 for MNIST; 5e − 4 for the others • Number of training steps: 2000 for CelebA-HQ; 1000 for the others
For Langevin Dynamics
• Learning rate: 5e − 4 for all datasets • Length of chain: 1000 for CIFAR-10; 2000 for the others
For PL-
MCMC • Learning rate: 5e −• Length of chain: 2000 for MNIST • σ a = 1e − 3 , σ p = 0 .
D.4 Hyperparameters: Inverse Problems
[Table]
Dataset
[Table]
[Table]
[Figure]