L OCALISED G ENERATIVE F LOWS
Anonymous authors
Paper under double-blind review
A BSTRACT
We argue that ﬂow-based density models based on continuous bijections are limited in their ability to learn target distributions with complicated topologies, and propose localised generative ﬂows (LGFs) to address this problem.
LGFs are composed of stacked continuous mixtures of bijections, which enables each bijection to learn a local region of the target rather than its entirety.
Our method is a generalisation of existing ﬂow-based methods, which can be used without modiﬁcation as the basis for an LGF model.
Unlike normalising ﬂows, LGFs do not permit exact computation of log likelihoods, but we propose a simple variational scheme that performs well in practice.
We show empirically that LGFs yield improved performance across a variety of density estimation tasks.
I NTRODUCTION
Flow-based generative models, often referred to as normalising ﬂows , have become popular methods for density estimation because of their ﬂexibility, expressiveness, and tractable likelihoods.
Given the problem of learning an unknown target density p (cid:63)X on a data space X , normalising ﬂows model p (cid:63)X as the marginal of X obtained by the generative process [Equation]
[Equation]
where p Z is a prior density on a space Z , and g :
X → Z is a bijection.
1 Assuming sufﬁcient regularity, it follows that X has density p X ( x ) =
p Z ( g ( x )) | det Dg ( x ) | (see e.g. Billingsley (2008)).
The parameters of g can be learned via maximum likelihood given i.i.d.
samples from p (cid:63)X .
To be effective, a normalising ﬂow model must specify an expressive family of bijections with tractable Jacobians.
Afﬁne coupling layers (Dinh et al., 2014; 2016), autoregressive transformations (Germain et al., 2015; Papamakarios et al., 2017), ODE-based transformations (Grathwohl et al., 2018), and invertible ResNet blocks (Behrmann et al., 2019) are all examples of such bijections that can be composed to produce complicated ﬂows.
These models have demonstrated signiﬁcant promise in their ability to model complex datasets (Papamakarios et al., 2017) and to synthesise novel data points (Kingma & Dhariwal, 2018).
However, in all these cases, g is continuous in x .
We believe this is a signiﬁcant limitation of these models since it imposes a global constraint on g − 1 , which must learn to match the topology of Z , which is usually quite simple, to the topology of X , which we expect to be very complicated.
We argue that this constraint makes maximum likelihood estimation extremely difﬁcult in general, leading to training instabilities and erroneous regions of high likelihood in the learned density landscape.
To address this problem we introduce localised generative ﬂows (LGFs), which generalise equationby replacing the single bijection g with stacked continuous mixtures of bijections { G ( · ; u ) } u ∈U for an index set U .
Intuitively, LGFs allow each G ( · ; u ) to focus on modelling only a local component of the target that may have a much simpler topology than the full density.
LGFs do not stipulate the form of G , and indeed any standard choice of g can be used as the basis of its deﬁnition.
We pay a price for these beneﬁts in that we can no longer compute the likelihood of our model exactly and must instead resort to a variational approximation, with our training objective replaced by the evidence lower bound (ELBO).
However, in practice we ﬁnd this is not a signiﬁcant limitation, as the bijective structure of LGFs permits learning a high-quality variational distribution suitable for large-scale training. We show empirically that LGFs outperform their counterpart normalising ﬂows across a variety of density estimation tasks.
L IMITATIONS OF N ORMALISING F LOWS
Consider a normalising ﬂow model deﬁned by a family of bijections g θ parameterised by θ .
Suppose we are in the typical case that each g θ is continuous in x .
Intuitively, this seems to pose a problem when p (cid:63)X and p Z are supported 2 on domains with different topologies, since continuous bijections (with continuous inverses) necessarily preserve topology.
This is made precise in Propositionbelow; we defer the proof to Appendix A.1.
In practice, suggestive pathologies along these lines are readily apparent in simple 2-D experiments as shown in Figure 1.
In this case, the density model (a masked autoregressive ﬂow (MAF) (Papamakarios et al., 2017)) is unable to continuously transform the support of the prior (a standard 2-D Gaussian) into the support of p (cid:63)X , which is the union of two disjoint rectangles and hence clearly has a different topology. {+Proposition 1 (Limitations of Normalising Flows) . Suppose each g θ is a homeomorphism and det Dg θ ( x ) (cid:54) = 
 for all θ and x .+}
If the prior p Z has a support that is not homeomorphic to p ∗ X , then no choice of θ will yield a normalising ﬂow model that matches p ∗ X exactly.
In other words, once we ﬁx p Z , we immediately restrict the set of target distributions that we may possibly express using our model. Provided each g θ is continuous, this occurs no matter how complicated we make our parameterisation – it holds even if we allow all continuous bijections. We believe this fact signiﬁcantly limits the ability of normalising ﬂows to learn any target distribution whose support has a complicated topology. Furthermore, note that the standard maximum likelihood objective is asymptotically equivalent to
[Equation]
where p θX ( x ) =
p
Z ( g θ ( x )) | det Dg θ ( x ) | and D KL denotes the Kullback-Leibler divergence.
We conjecture that the problem identiﬁed in Proposition 1 is further exacerbated by this objective.
Note that equation 2 is inﬁnite unless supp p θX ⊇ supp p (cid:63)X , but strictly positive unless supp p θX = supp p (cid:63)X .
Thus equation 2 encourages the support of p θX to approximate that of p (cid:63)X from above as closely as possible, but any overshoot – however small – carries an immediate inﬁnite penalty.
This seems problematic for gradient-based methods of solving equation 2, especially when, in light of
Proposition 1, we expect an optimal θ to be pathological, since g θ is encouraged to approximate some function that is not a continuous bijection. In this context, for near-optimal θ , it seems plausible that supp p θX might be very sensitive to perturbations of θ , leading to a highly-varying loss landscape when near the optimum.
We conjecture this makes the landscape around the optimal θ difﬁcult to navigate for gradient-based methods, leading to increasing ﬂuctuations and instabilities strictly positive.
However, our statements are also approximately true when supp p is interpreted as the region of p which is not smaller than some threshold.
This is relevant in practice since, even if both are highly concentrated on some small region, it is common to assume that p (cid:63)X and p Z have full support, in which case the supports of p (cid:63)X and p Z would be trivially homeomorphic.
of the KL
in equation 2 that degrade the quality of the ﬁnal g θ produced.
Figure 2 illustrates this in practice: observe that, even after 300 training epochs, the support of the standard MAF is unstable.
A simple way to ﬁx these problems would be to use a more complicated p Z that is better matched to the structure of p (cid:63)X .
For instance, taking p Z to be a mixture model has previously been found to improve the performance of normalising ﬂows in some cases (Papamakarios et al., 2017).
However, this approach requires prior knowledge of the topology of supp p (cid:63)X that might be difﬁcult to obtain: e.g. solving the problem with a Gaussian mixture p Z requires us to know the number of connected components of supp p (cid:63)X beforehand, and even then would require that these components are each homeomorphic to a hypersphere.
Ideally, we would like our model to be ﬂexible enough to learn the structure of the target on its own, with minimal explicit design choices required on our part.
An alternative approach would be to try a more expressive family g θ in the hope that this better conditions the optimisation problem in equation 2.
Several works have considered families of g θ that are (in principle) universal approximators of any continuous probability densities (Huang et al., 2018;
Jaini et al., 2019).
While we have not performed a thorough empirical evaluation of these methods, we suspect that these models can at best mitigate the problems described above, since the assumption of continuity of g θ in practice holds for universal approximators also.
Moreover, the method we propose below can be used in conjunction with any standard ﬂow, so that we expect even better performance when an expressive g θ is combined with our approach.
Finally, we note that the shortcomings of the KL divergence for generative modelling are described at length by Arjovsky et al.
There the authors suggest instead using the Wasserstein distance to measure the discrepancy between p (cid:63) X and p Z , since under typical assumptions this will yield a continuous loss function suitable for gradient-based training.
However, the Wasserstein distance is difﬁcult to estimate in high dimensions, and its performance can be sensitive to the choice of ground metric used (Peyr´e et al., 2019).
Our proposal here is to keep the KL objective in equation 2 and instead modify the model so that we are not required to map supp p Z onto supp p (cid:63)X using a single continuous bijection.
We describe our method in full now.
L OCALISED GENERATIVE FLOWS 3. M ODEL
The essence of our idea is to replace the single g used in equation 2 with an indexed family { G ( · ; u ) } u ∈U such that each G ( · ; u ) is a bijection from X to Z .
Intuitively, our aim is for each
G ( · ; u ) only to push the prior onto a local region of supp p (cid:63) X , thereby relaxing the constraints posed by standard normalising ﬂows as described above.
To do so, we now deﬁne p X as the marginal density of X obtained via the following generative process:
[Equation]
Here p U | Z is an additional term that we must specify.
In all our experiments we take this to be a mean ﬁeld Gaussian, so that p U | Z ( u | z ) =
Gaussian( u ; µ ( z ) , σ ( z ) 2 I d u ) for some functions µ, σ : [Equation]
[Figure] Z → U ⊆ R d u , where d u is the dimension of the index set U , and I d u is the d u × d u identity matrix.
Other possibilities, such as the Concrete distribution (Maddison et al., 2016), might also be useful.
Informally, 3 this yields the joint model p X,U,Z ( x, u, z ) =
p Z ( z ) p U | Z ( u | z ) δ ( x − G − 1 ( z ; u )) , where δ is the Dirac delta.
We can marginalise out the dependence on z by making the change of variable z = G ( x (cid:48) ; u ) , which means dz = | det DG ( x (cid:48) ; u ) | dx (cid:48) . This yields a proper density for ( X, U ) via
[Equation]
We then obtain our density model p X by integrating over u :
[Equation]
In other words, p X ( x ) is a mixture (in general, inﬁnite) of individual normalising ﬂows G ( x ; u ) , each weighted by p U | Z ( u | G ( x ; u )) .
We can also stack this architecture by taking p Z itself to be a density of the form of equation 4.
Doing so with L layers of stacking corresponds to the marginal of X ≡ Z L obtained via:
[Equation]
where now each G (cid:96) ( · ; u ) :
X → Z is a bijection for all u ∈ U .
The stochastic computation graph corresponding to this model is shown in Figure 3a.
In this case, the same argument yields
[Equation]
where p Z 0 ,U 1:0 ≡ p Z 0 .
This approach is in keeping with the standard practice of constructing normalising ﬂows as the composition of simpler bijections, which can indeed be recovered here by taking each p U (cid:96) | Z (cid:96) − 1 ( ·| Z (cid:96) − 1 ) to be Dirac.
We have found stacking to improve signiﬁcantly the overall expressiveness of our models, and make extensive use of it in our experiments below.
3. B ENEFITS
Heuristically, we believe our model allows each G ( · ; u ) to learn a local region of p (cid:63)X , thereby greatly relaxing the global constraints on existing ﬂow-based models described above.
To ensure a ﬁnite KL, we no longer require the density
p
Z
( G ( x ; u )) | det DG ( x ; u ) | to have support covering the entirety of supp p (cid:63) X for any given u ; all that matters is that every region of supp p (cid:63) X is covered for some u .
Our model can thus achieve good performance with each bijection G ( · ; u ) faithfully mapping onto a potentially very small component of p (cid:63)X .
This seems intuitively more achievable than the previous case, wherein a single bijection is required to capture the entire target.
This argument is potentially clearest if u is discrete.
For example, even if u ∈ { u − 1 , u 1 } can take only two possible values, it immediately becomes simpler to represent the target shown in Figure 1 using a Gaussian prior: we simply require G − 1 ( · ; u − 1 ) to map onto one component, and G − 1 ( · ; u 1 ) to map onto the other.
In practice, we could easily implement such a G using two separate normalising ﬂows that are trained jointly.
The discrete case is also appealing since in principle it allows exact evaluation of the integral in equation 4, which becomes a summation.
Unfortunately this approach also has signiﬁcant drawbacks that we discuss at greater length in Appendix B.
We therefore instead focus here on a continuous u .
In this case, for example, we can recover p (cid:63)X from Figure 1 by partitioning U into disjoint regions U − 1 and U 1 , and having G − 1 ( · ; u ) map onto the left component of p (cid:63)X for u ∈ U − 1 , and the right component for u ∈ U 1 .
Observe that in this scenario we do not require any given G − 1 ( · ; u ) to map onto both components of the target, which is in keeping with our goal of localising the model of p (cid:63)X that is learned by our method.
In practice G will invariably be continuous in both its arguments, in which case it will not be possible to partition U disjointly in this way.
Instead we will necessarily obtain some additional intermediate region U 0 on which G − 1 ( · ; u ) maps part of supp p
Z outside of supp p (cid:63)X , so that p X ( x ) will be strictly positive there.
However, Proposition 2 below shows (proof in Appendix A.2) that this is not a problem for LGFs – we are able to use p U | Z to downweight such a region, which avoids the support mismatch issue of standard normalising ﬂows. Proposition 2 (Beneﬁt of Localised Generative Flows) .
Suppose that supp p ∗ X is open and G ( · ; u ) is continuous for each u . Suppose further that, for each x ∈ supp p ∗ X , the set [Equation]
has positive Lebesgue measure.
Then there exists p U | Z such that supp p X = supp p [Equation]
X .
3. I NFERENCE
Even in the single layer case ( L = 1 ), if u is continuous, then the integral in equation 4 is intractable.
In order to train our model, we therefore resort to a variational approximation: we introduce an approximate posterior q U 1: L | X ≈ p U 1: L | X , and consider the evidence lower bound (ELBO) of log p X ( x ) :
[Equation]
X ( u 1: L | x ) (cid:3) . It is straightforward to show that L ( x ) ≤ log p X ( x ) , and that this bound is tight when q U 1: L | X is the exact posterior p U 1: L | X . This allows learning an approximation to p (cid:63)X by maximising n − 1 (cid:80) ni
=1 L ( x i ) jointly in p X,U 1: L and q U 1: L | X , for a dataset of n i.i.d.
samples x i ∼ p (cid:63)X .
It can be shown (see Appendix A.3) that the exact posterior factors as p U 1: L | X ( u 1: L | x ) = (cid:81)
L(cid:96)
= [Equation] p [Equation]
(cid:96) | Z (cid:96) ( u (cid:96) | z (cid:96) ) , where z L := x , and z (cid:96) − 1 := G (cid:96) ( z (cid:96) ; u (cid:96) ) for (cid:96) ≤ L .
We thus endow q U 1: L | X with the same form:
[Equation]
The stochastic computation graph for this inference model is shown in Figure 3b.
In conjunction with equation 5, this allows writing the ELBO recursively as
[Equation]
; u (cid:96) ) | − log q U (cid:96) | Z (cid:96) ( u (cid:96)
| z (cid:96) ) (cid:3) for (cid:96) ≥ 1 , with the base case L 0 ( z 0 ) =
p Z 0 ( z 0 ) .
Here we recover L ( x ) ≡ L L ( z L ) .
Now let θ denote all the parameters of both p X,U 1: L and q U 1: L | X .
Suppose each q U (cid:96) | Z (cid:96) can be suitably reparametrised (Kingma & Welling, 2013; Rezende et al., 2014)
so that h (cid:96) ( (cid:15) (cid:96) , z (cid:96) ) ∼ q U (cid:96) | Z (cid:96) ( ·| z (cid:96) ) when (cid:15) (cid:96) ∼ η (cid:96) for some function h (cid:96) and density η (cid:96) , where η (cid:96) does not depend on θ .
In all our experiments we give q U (cid:96) | Z (cid:96) the same mean ﬁeld form as in equation 3, in which case this holds immediately as described e.g. by Kingma & Welling (2013).
We can then obtain unbiased estimates of ∇ θ L ( x ) straightforwardly using Algorithm 1, which in turn allows minimising our objective via stochastic gradient descent.
Note that although this algorithm is speciﬁed in terms of a single value of z (cid:96) , it is trivial to obtain an unbiased estimate of ∇ θ m − 1 (cid:80) mj
=1 L ( x j ) for a minibatch of points { x j } mj =1 by averaging over the batch index at each layer of recursion.
Algorithm
RAD ELBO( z (cid:96) , (cid:96) ) if (cid:96) = 0 then return ∇ θ log p
[Equation]
z
[Equation] ) else [Equation] u
[Equation]
∆
[Equation]
) return G RAD ELBO( z (cid:96) − 1 , (cid:96) − 1 ) + ∇ θ ∆ (cid:96)
3.3. P ERFORMANCE
A major reason for the popularity of normalising ﬂows is the tractability of their exact log likelihoods.
In contrast, the variational scheme described here can produce at best an approximation to this value, which we might expect reduces performance of the ﬁnal density estimator learned.
Moreover, particularly when the number of layers L is large, it might seem that the variance of gradient estimators obtained from Algorithm 1 would be impractically high.
[-(cid:96) function G RAD ELBO( z (cid:96) , (cid:96) ) if (cid:96) = 0 then return p Z 
 ( z (cid:96) ) else (cid:15) [Table]-]However, in practice we have not found either of these problems to be a signiﬁcant limitation, as our experimental results in Section 5 show.
Empirically we ﬁnd improved performance over standard ﬂows even when using the ELBO as our training objective.
We also ﬁnd that importance sampling is sufﬁcient for obtaining good, low-variance (if slightly negatively biased) estimates of log p X ( x ) (Rezende et al., 2014, Appendix E) at test time, although we do note that the stochasticity here can lead to small artefacts like the occasional white spots visible above in our 2-D experiments.
We similarly do not ﬁnd that the variance of Algorithm 1 grows intractably when we use a large number of layers L , and in practice we are able to train models having the same depth as popular normalising ﬂows.
We conjecture that this occurs because, as the number of stacked bijections in our generative process grows, the complexity required of each individual bijection to map p Z to p (cid:63) X naturally decreases.
We therefore have reason to think that, as L becomes large, learning each q U (cid:96) | Z (cid:96) will become easier, so that the variance at each layer will decrease and the overall variance will remain fairly stable.
3. C HOICE OF INDEXED BIJECTION FAMILY
We now consider the choice of G , for which there are many possibilities.
In our experiments, which all take Z = R d , we focus on the simple case of
[Equation]
where s, t : U → Z are unrestricted mappings, g : X → Z is a bijection, and (cid:12) denotes elementwise multiplication.
In this case log | det DG ( x ; u ) | = log | det Dg ( x ) | + (cid:80) di
{+=
+} [ s ( u )] i , where [ s ( u )]
i is the i th component of s ( u ) .
This has the advantage of working out-of-the-box with all pre-existing normalising ﬂow methods for which a tractable Jacobian of g is available.
We provide alternative suggestions to this construction in Appendix C.
Equation 6 also has an appealing similarity with the common practice of applying afﬁne transformations between ﬂow steps for normalisation purposes, which has been found empirically to improve stability, convergence time, and overall performance (Dinh et al., 2016;
Papamakarios et al., 2017;
Kingma & Dhariwal, 2018).
In prior work, s and t have been simple parameters that are learned either directly as part of the model, or updated according to running batch statistics.
Our approach may be understood as a generalisation of these techniques.
R ELATED W ORK 4. D ISCRETE M IXTURE M ETHODS
Several density models exist that involve discrete mixtures of normalising ﬂows.
A closely-related approach to ours is RAD (Dinh et al., 2019), which is a special case of our model when L = 1 , u is discrete, and X is partitioned disjointly.
In the context of Monte Carlo estimation, Duan (2019) proposes a similar model to RAD that does not use partitioning.
Ziegler & Rush (2019) introduce a normalising ﬂow model for sequences with an additional latent variable indicating sequence length.
More generally, our method may be considered an addition to the class of deep mixture models (Tang et al., 2012; Van den Oord & Schrauwen, 2014), with our use of continuous mixing variables designed to reduce the computational difﬁculties that arise when stacking discrete mixtures hierarchically – we refer the reader to Appendix B for more details on this.
4. M ETHODS C OMBINING V ARIATIONAL I NFERENCE AND N ORMALISING F LOWS
There is also a large class of methods which use normalising ﬂows to improve the inference procedure in variational inference (van den Berg et al., 2018; Kingma et al., 2016; Rezende & Mohamed, 2015), although ﬂows are not typically present in the generative process here.
This approach can be described as using normalising ﬂows to improve variational inference , which contrasts our goal of using variational inference to improve normalising ﬂows .
However, there are indeed some methods augmenting normalising ﬂows with variational inference, but in all cases below the variational structure is not stacked to obtain extra expressiveness.
Ho et al.
(2019) use a variational scheme to improve upon the standard dequantisation method for deep generative modelling of images (Theis et al., 2015); this approach is orthogonal to ours and could indeed be used alongside LGFs.
Gritsenko et al.
(2019) also generalise normalising ﬂows using variational methods, but they incorporate the extra latent noise into the model in a much more restrictive way.
Das et al.
(2019) only learn a low-dimensional prior over the noise space variationally.
4. S TACKED V ARIATIONAL M ETHODS
Finally, we can contrast our approach with purely variational methods which are not ﬂow-based, but still involve some type of stacking architecture.
The main difference between these methods and LGFs is that the bijections we use provide us with a generative model with far more structure, which allows us to build appropriate inference models more easily.
Contrast this with, for example, Rezende et al.
(2014), in which the layers are independently inferred, or Sønderby et al.
(2016), which requires a complicated parameter-sharing scheme to reliably perform inference.
A singlelayer instance of our model also shares some similarity with the fully-unsupervised case in Maaløe et al.
, but the generative process there conditions the auxiliary variable on the data.
E XPERIMENTS
We evaluated the performance of LGFs on several problems of varying difﬁculty, including synthetic 2-D data, several UCI datasets, and two image datasets. We describe the results of the UCI and image datasets here; Appendix D.2 contains the details of our 2-D experiments.
In each case, we compared a baseline ﬂow to its extension as an LGF model with roughly the same number of parameters.
We obtained each LGF by inserting a mixing variable U (cid:96) between every component layer of the baseline ﬂow. For example, we inserted a U (cid:96) after each autoregressive layer in MAF (Papamakarios et al., 2017), and after each afﬁne coupling layer in RealNVP (Dinh et al., 2016).
In all cases we obtained p U (cid:96) | Z (cid:96) − 1 , q U (cid:96) | Z (cid:96) , s (cid:96) , and t (cid:96) as described in Appendix D.1.
We inserted batch normalisation (Ioffe & Szegedy, 2015) between ﬂow steps for the baselines as suggested by Dinh et al.
(2016), but omit it for the LGFs, since our choice of bijection family is a generalisation of batch normalisation as described in Subsection 3.4.
We trained our models to maximise either the log-likelihood (for the baselines) or the ELBO (for the LGFs) using the ADAM optimiser (Kingma & Ba, 2014) with default hyperparameters and no weight decay.
For the UCI and image experiments we stopped training after 30 epochs of no validation improvement for the UCI experiments, and after 50 epochs for the image experiments.
Both validation and test performance were evaluated using the exact log-likelihood for the baseline, and the standard importance sampling estimator of the average log-likelihood (Rezende et al., 2014,
Appendix E) for LGFs.
For validation we used 5 importance samples for the UCI datasets andfor the image datasets, while for testing we used 1000 importance samples in all cases.
Our code is available at https://github.com/anonsubmission974/lgf .
5. UCI D ATASETS
We tested the performance of LGFs on the POWER, GAS, HEPMASS, and MINIBOONE datasets from the UCI repository (Bache & Lichman, 2013).
We preprocessed these datasets identically to Papamakarios et al.
(2017), and use the same train/validation/test splits.
For all models we used a batch size of 1000 and a learning rate of 10 − 3 . These constitute a factor of 10 increase over the values used by Papamakarios et al. and were chosen to decrease training time.
Our baseline results differ slightly from Papamakarios et al. (2017), which may be as a result of this change.
We focused on MAF as our baseline normalising ﬂow because of its improved performance over alternatives such as RealNVP for general-purpose density estimation (Papamakarios et al., 2017).
A given MAF model is deﬁned by how many autoregressive layers it uses as well as the sizes of the autoregressive networks at each layer.
For an LGF-MAF, we must additionally deﬁne the neural networks used in the generative and inference process, for which we use multi-layer perceptrons (MLPs).
We considered a variety of choices of hyperparameters for MAF and LGF-MAF.
Each instance of LGF-MAF had a corresponding MAF conﬁguration, but in order to compensate for the parameters introduced by our additional neural networks, we also considered deeper and wider MAF models. For each LGF-MAF, the dimensionality of u (cid:96) was roughly a quarter of that of the data. Full hyperparameter details are given in Appendix D.3.
For three random seeds, we trained models using [Table] every hyperparameter conﬁguration, and then chose the best-performing model across all parameter conﬁgurations using validation performance.
Table 1 shows the resulting test-set log-likelihoods averaged across the different seeds.
It is clear that LGF-MAFs yield improved results in this case.
5. I MAGE DATASETS
We also considered LGFs applied to the Fashion-MNIST (Xiao et al., 2017) and CIFAR- (Krizhevsky et al., 2009) datasets.
In both cases we applied the dequantisation scheme of Theis et al.
(2015) beforehand, and trained all models with a learning rate of 10 − 4 and a batch size of 100.
We took our baseline to be a RealNVP with the same architecture used by Dinh et al.
(2016) for their CIFAR-10 experiments. In particular, we used 10 afﬁne coupling layers with the corresponding alternating channelwise and checkerboard masks. Each coupling layer used a ResNet (He et al., 2016a;b) consisting of 8 residual blocks of 64 channels (denoted 8 × 64 for brevity).
We also replicated their multi-scale architecture, squeezing the channel dimension after the ﬁrst 3 coupling layers, and splitting off half the dimensions after the ﬁrst 6.
This model had 5.94M parameters for Fashion-MNIST and 6.01M parameters for CIFAR-10.
For completeness, we also consider a RealNVP model with coupler networks of size 4 × 64 to match those used below in LGF-RealNVP. This model had 2 .
99 M and 3 .
05 M parameters for Fashion-MNIST and CIFAR-10 respectively.
error of test-set bits per dimension. RealNVP ( k ) refers to a RealNVP model with k residual blocks in the coupling networks. [Table] For the LGF-RealNVP, we sought to maintain roughly the same depth over which we propagate gradients as in the baseline.
To this end, our coupling networks were ResNets of size 4 × 64 , and each p U (cid:96) | Z (cid:96) − 1 and q U (cid:96) | Z (cid:96) used ResNets of size 2 × 64 .
Our ( s (cid:96) , t (cid:96) ) network was a ResNet of size 2 × 8 .
We give u (cid:96) the same shape as a single channel of z (cid:96) , and upsampled to the dimension of z (cid:96) by adding channels at the output of the ( s (cid:96) , t (cid:96) ) ResNet.
Our model had 5.99M parameters for Fashion-MNIST and 6.07M parameters for CIFAR-10.
Table 2 shows that LGFs consistently outperformed the baseline models.
We moreover found that
LGFs tend to train faster : the average epoch with best validation performance on CIFAR-10 wasfor LGF-RealNVP, and 723 for RealNVP (8).
5 Samples from all models are shown in Appendix D.4.
We also found that using the ELBO instead of the log-likelihood does not penalise our method.
The gap between the estimated test-set log-likelihood and the average test-set ELBO was not very large for the LGF models, with a relative error of 8 .
98 × 10 − 3 for Fashion-MNIST and 6 . 88 × 10 −for CIFAR-10.
Moreover, the importance-sampling-based log-likelihood estimator itself had very low variance when using 1000 importance samples.
For each trained LGF model, we estimated the relative standard deviation using three separate samples of this estimator. We obtained an average over all trained models of 8 .
34 × 10 − 4 for Fashion-MNIST, and 2 .
07 × 10 − 5 for CIFAR-10.
C ONCLUSION
In this paper, we have proposed localised generative ﬂows for density estimation, which generalise existing normalising ﬂow methods and address the limitations of these models in expressing complicated target distributions.
Our method obtains successful empirical results on a variety of tasks.
Many extensions appear possible, and we believe localised generative ﬂows show promise as a means for improving the performance of density estimators in practice.
A A DDITIONAL P ROOFS A. P ROOF OF P ROPOSITION 1 (S EE PAGE 2)
{+Proposition 1 (Limitations of Normalising Flows) . Suppose each g θ is a homeomorphism and det Dg θ ( x ) (cid:54) = 
 for all θ and x .+}
If the prior p Z has a support that is not homeomorphic to p ∗ X , then no choice of θ will yield a normalising ﬂow model that matches p ∗ X exactly.
Proof.
For any choice of θ , let p θX ( x ) be the normalising ﬂow model produced by p Z and g θ , i.e.
[Equation]
We can bound the total variation distance d TV between p θX and p (cid:63)X by
[Equation]
Observe that this is strictly positive unless supp p (cid:63) X = supp p θ X (up to a set of p θ X -probability 0).
But since g θ is a homeomorphism, this is not possible.
Hence d TV ( p (cid:63)X , p θX ) > 0 , so that p (cid:63)X (cid:54) = p θX .
A. P ROOF OF P ROPOSITION
2 (S EE PAGE 5)
Proposition 2 (Beneﬁt of Localised Generative Flows) .
Suppose that supp p ∗ X is open and G ( · ; u ) is continuous for each u .
Suppose further that, for each x ∈ supp p ∗ X , the set
[Equation]
has positive Lebesgue measure.
Then there exists p U | Z such that supp p X = supp p ∗ X .
Proof.
Take p
U
to be any positive density.
Deﬁne
[Equation]
Observe that each A u has positive Lebesgue measure, since it is the continuous injective image of an open set.
Thus, for each u , we can choose p Z | U ( ·| u ) to be any positive density on A u .
This allows us to deﬁne
[Equation]
We claim this construction gives supp p X = supp p ∗ X .
First suppose x ∈ supp p ∗ X .
For each u , observe that p Z | U ( G ( x ; u ) | u ) > 0 by the deﬁnition of A u , and hence p U | Z ( u | G ( x ; u )) > 0 since p U is positive.
Consequently,
[Equation]
since the integrand is strictly positive on B x , which has positive measure.
Now suppose x (cid:54)∈ supp p ∗ X .
For each u , since p Z | U ( ·| u ) is supported on A u , we must have p Z |
U
( G ( x ; u ) | u ) = 0 , and hence p U | Z ( u | G ( x ; u ))
= 0 .
From this it follows directly that p X ( x ) = 0 , i.e. x (cid:54)∈ supp p X .
A. C ORRECTNESS OF P OSTERIOR F ACTORISATION
We want to prove that the exact posterior factorises according to
[Equation]
where Z L ≡ X , z L ≡ x , and z (cid:96) − 1 := G (cid:96) ( z (cid:96) ; u (cid:96) ) for (cid:96) ≤ L .
It is convenient to introduce the index notation U >(cid:96) ≡ U (cid:96) +1: L .
First note that we can always write this posterior autoregressively as
[Equation]
Now, consider the graph of the full generative model shown in Figure 3a.
It is clear that all paths from U (cid:96) to any node in the set H (cid:96) :
= { U >(cid:96) , Z >(cid:96) } are blocked by Z (cid:96) , and hence that Z (cid:96) d-separates H (cid:96) from U (cid:96) (Bishop, 2006).
Consequently,
[Equation]
Informally, 6 we then have the following for any (cid:96) ∈ { 1 , . . .
, L − 1 } :
[Equation]
=
[Equation]
by the deﬁnition of z (cid:96) , where we write G (cid:96) ( · ; u ) ≡ G (cid:96),u ( · ) to remove any ambiguities when composing functions.
We can substitute this result into equation 8 to obtain equation 7.
B D ISCRETE u
While a discrete u is appealing for its ability to produce exact likelihoods, it also suffers several disadvantages that we describe now.
First, observe that the choice of the number of discrete values taken by u has immediate implications for the number of disconnected components of supp p (cid:63)X that G can separate, which therefore seems to require making fairly concrete assumptions (perhaps implicitly) about the topology of the target of interest.
To mitigate this, we might try taking the number of u values to be very large, but then in turn the time required to evaluate equation 4 (now a sum rather than an integral) necessarily increases.
This is particularly true when using a stacked architecture, since to evaluate p X ( x ) with L layers each having K possible u -values takes Θ( K L ) complexity.
Dinh et al.
(2019) propose a model that partitions X so that only one component in each summation is nonzero for any given x , which reduces this cost to Θ( L ) .
However, this partitioning means that their p X is not continuous as a function of x , which is reported to make the optimisation problem in equation 2 difﬁcult.
Unlike for continuous u , the ELBO objective is also of limited effectiveness in the discrete case.
Since p U (cid:96) | Z (cid:96) − 1 deﬁnes a discrete distribution at each layer, we would need a discrete variational approximation q [Table] to ensure a well-deﬁned ELBO. However, the parameters of a discrete distribution are not ameanble to the reparametrisation trick, and hence we would expect our gradient estimates of the ELBO to have high variance.
As mentioned above, a compromise here might be to use the CONCRETE distribution (Maddison et al., 2016) to approximate a discrete u while still apply variational methods. We leave exploring this for future work. C O THER C HOICES OF I NDEXED B IJECTION F AMILY Other choices of G than our suggestion in Subsection 3.4 are certainly possible.
For instance, we have had preliminary experimental success on small problems by simply taking g to be the identity, in which case the model is greatly simpliﬁed by not requiring any Jacobians at all. Alternatively, it is also frequently possible to modify the architecture of standard choices of g to obtain an appropriate G . For instance, afﬁne coupling layers, a key component of models such as RealNVP (Dinh et al., 2016), make use of neural networks that take as input a subset of the dimensions of x .
By concatenating u to this input, we straightforwardly obtain a family of bijections G ( · ; u ) for each value of u .
This requires more work to implement than our suggested method, but has the advantage of no longer requiring a choice of s and t .
We have again had preliminary empirical success with this approach.
We leave a more thorough exploration of these alternative possibilites for future work.
D F URTHER E XPERIMENTAL D ETAILS D. LGF ARCHITECTURE
In addition to the bijections from the baseline ﬂow, an LGF model of the form we consider requires specifying p U (cid:96) | Z (cid:96) − 1 , q U (cid:96) | Z (cid:96) , s (cid:96) and t (cid:96) .
In all our experiments
• p U (cid:96) | Z (cid:96) − 1 ( ·| z (cid:96) − 1 ) =
Normal (cid:0) µ p ( z (cid:96) − 1 ) , σ p ( z (cid:96) − 1 ) 2 (cid:1) , where µ p and σ p were two separate outputs of the same neural network • q U (cid:96)
| Z (cid:96) ( ·| z (cid:96) ) =
Normal (cid:0) µ q
( z (cid:96) ) , σ q ( z (cid:96) ) 2 (cid:1) , where µ q and σ q were two separate outputs of the same neural network • s (cid:96) and
t (cid:96) were two separate outputs of the same neural network.
To gain intuition about our model, we ran several experiments on the simple 2-D datasets shown in Figure 1 and Figure 4.
Speciﬁcally, we compared the performance of a baseline MAF against an
LGF-MAF.
We used MLPs for all networks involved in our model.
For the dataset shown in Figure 1, the baseline MAF had 10 autoregressive layers consisting of 4 hidden layers with 50 hidden units.
The LGF-MAF had 5 autoregressive layers consisting ofhidden layers with 10 hidden units.
Each u (cid:96) was a single scalar. The ( s (cid:96) , t (cid:96) ) network consisted ofhidden layers of 10 hidden units, and both the ( µ p , σ p ) and ( µ q , σ q ) networks consisted of 4 hidden layers of 50 hidden units.
In total the baseline MAF had 80080 parameters, while LGF-MAF had 80810 parameters.
We trained both models for 300 epochs.
We used more parameters for the datasets shown in Figure 4, since these targets have more complicated topologies.
In particular, the baseline MAF had 20 autoregressive layers, each with the same structure as before.
The LGF-MAF had 5 autoregressive layers, now with 4 hidden layers ofhidden units.
The other networks were the same as before, and each u (cid:96) was still a scalar.
In total the baseline MAF had 160160 parameters, while our model had 119910 parameters.
We trained all models now for 500 epochs.
The results of these experiments are shown in Figure 1 and Figure 4.
Observe that LGF-MAF consistently produces a more faithful representation of the target distribution than the baseline.
A failure mode of our approach is exhibited in the spiral dataset, where our model still lacks the power to fully capture the topology of the target.
However, we did not ﬁnd it difﬁcult to improve on this: by increasing the size of the mean/standard deviation networks to 8 hidden layers of 50 hidden units (and keeping all other parameters ﬁxed), we were able to obtain the result shown in Figure 5.
This model had a total of 221910 parameters.
For the sake of a fair comparison, we also tried increasing the complexity of the MAF model, by the size of its autoregressive networks to 8 hidden layers of 50 hidden units (obtaining 364160 parameters total).
This model diverged after approximatelyepochs.
The result after 150 epochs is shown in Figure 5.
D. UCI E XPERIMENTS
In Tables 3, 4, and 5, we list the choices of parameters for MAF and LGF-MAF.
In all cases, we allowed the base MAF to have more layers and deeper coupler networks to compensate for the additional parameters added by the additional components of our model.
Note that neural networks are listed as size K 1 × K 2 , where K 1 denotes the number of hidden layers and K 2 denotes the size of the hidden layers.
All combinations of parameters were considered; in each case, there wereconﬁgurations for MAF and 8 conﬁgurations for LGF-MAF [Table] [Table]
[Table] [Table]
D. I MAGE E XPERIMENTS
In ﬁgures 6 to
[Table]
we present some samples synthesised from the density models trained on FashionMNIST and CIFAR-
[Table]