Latent Space Reﬁnement for Deep Generative Models
Abstract
Deep generative models are becoming widely used across science and industry for a variety of purposes.
A common challenge is achieving a precise implicit or explicit representation of the data probability density.
Recent proposals have suggested using classiﬁer weights to reﬁne the learned density of deep generative models.
We extend this idea to all types of generative models and show how latent space reﬁnement via iterated generative modeling can circumvent topological obstructions and improve precision.
This methodology also applies to cases were the target model is non-differentiable and has many internal latent dimensions which must be marginalized over before reﬁnement.
We demonstrate our Latent
Space Reﬁnement (LaSeR) protocol on a variety of examples, focusing on the combinations of Normalizing Flows and Generative Adversarial Networks.
We make all codes publicly available.
Introduction
Generative models are essential tools for many aspects of scientiﬁc and engineering workﬂows.
First-principles simulations encode physical laws and then samples from these simulations can be used for designing, performing, and interpreting a measurement.
However, these physics-based simulations can be too slow or not precise enough for a growing number of studies.
Deep learningbased techniques such as Generative Adversarial Networks (GAN) [1, 2], Variational Autoencoders (VAE) [3, 4] and Normalizing Flows (NF) [5, 6] are powerful surrogates that can accelerate slow simulations and model complex datasets that would otherwise be intractable to describe from ﬁrst principles.
For example, a growing number of studies are exploring these tools for high energy physics (HEP) applications [7–71].
A key difference in generative modeling between many scientiﬁc applications, including HEP, and typical industrial applications is that individual samples are often not useful.
Inference is performed on a statistical basis and so it is essential to model the probability density precisely and not just match its support.
Existing deep generative models have shown great promise in qualitatively modeling complex probability densities, but it is often challenging to achieve precision.
A useful strategy to improve the precision of generative models is to reﬁne their predictions.
For example, Ref.
[72] showed how the classiﬁcation part of a GAN can be used to reweight and resample the input probability distribution of the random noise.
In Ref.
[73] this classiﬁer information is already used during training to learn from an optimized latent space.
A similar idea was introduced in Ref.
[41] that is not speciﬁc to GANs, whereby a classiﬁer network is trained on the generated samples to produce weights that improve the precision of the generated probability distribution.
We combine and extent these approaches by introducing the La tent S pac e R eﬁnement (L A S E R) protocol, illustrated in Fig.
Like D CTR GAN [41, 74], L A S E R starts with any generative model g ( z ) and trains a post-hoc classiﬁer to distinguish the generated samples from the target data.
A challenge with D CTR GAN is that the results are weighted, which reduces the statistical power of a generated sample.
The energy-based framework of Discriminator Driven Latent Sampling (DDLS) [72] produces unweighted samples by transferring the weights to the latent space and then performing Langevin
Markov Chain Monte Carlo (MCMC) to produce unweighted samples.
We develop a more general protocol that works for any generative model and is more efﬁcient than the Langevin MCMC approach of Ref.
We begin by pulling back the weights from our post-hoc classiﬁer to the latent space, either directly or via the O MNI F OLD method [75] when the generator is not surjective.
We then propose to learn a second generative model Φ( y ) that maps an auxiliary latent space onto the reﬁned latent space.
Generating from the reﬁned model amounts to sampling from the auxiliary latent space and applying g (Φ( y )) .
We focus on the case where g is a normalizing ﬂow and Φ is a GAN because GANs need not to be invertible ; however, the method can be applied to any pair of generative models.
Furthermore, the procedure can be iterated for further reﬁnement.
Learning a post-hoc generative model in a latent space was also studied by the authors of Refs.
[76–78], where a standard autoencoder becomes probabilistic via a second model such as a NF trained on the bottleneck layer.
Similarly, the authors of Ref.
[79] proposed an iterative VAE setup to bring the latent space closer to a target multidimensional Gaussian random variable.
Further approaches in the context of
VAEs include iterative reﬁnement by Ref.
[80], iterative posterior inference by Refs.
[81, 82] as well as Ladder VAE [83] which introduce a simple form of posterior revision.
This paper is organized as follows.
We review related work and describe the statistical properties and challenges associated with existing generative models in Section 2.
We introduce various ways of implementing the L A S E R protocol in Sec. 3.
Illustrative numerical examples are provided in Sec.and the paper ends with conclusions and outlook in Sec. 5.
Background
A generator is a function g that maps a latent space Z ⊆ R N onto a target or feature space X ⊆ R M , with underlying probability densities p Z and p X , respectively.
Typically, p Z is chosen to be simple (e.g. normal or uniform) so that it is efﬁcient to generate data Z ∼ p
Z .
In some cases, the latent space
Z is broken into two components: one component that speciﬁes the input features of interest and one component that speciﬁes auxiliary latent features that are marginalized over.
When this happens, the map g is viewed as a stochastic function of the ﬁrst latent space component.
While our examples exclusively cover the case of deterministic maps from the full latent space to the target space, our approach can accommodate both settings as we explain in more detail below.
The function g can be constructed from ﬁrst-principles insights about the dynamics of a system or it can be learned from data.
For example, commonly-used deep generative models include
• Generative adversarial networks (GANs) [1, 2] • Variational autoencoders (VAEs) [3, 4] • Normalizing ﬂows (NFs)
Apart from the speciﬁc learning objective, all generative models have their intrinsic advantages and disadvantages.
The relevant features to understand the L A S E R protocol are summarized in Tableand are further discussed in the following.
Generative models and coordinate transformations
While the three generative models introduced in the previous section can all be trained directly on unlabeled data, they have different strategies for estimating p X .
GANs and VAEs learn this density implicitly by introducing an auxiliary task.
In the case of GANs, the auxiliary task is performed by a [Figure] [Table] discriminator network that tries to distinguish samples drawn from p Z passed through g and those drawn from p X directly.
For VAEs, the generator is called the decoder and the auxiliary task requires an encoder network h to satisfy g ( h ( x ))
≈ x , while regularizing the latent space probability density.
Due to the structure of these networks, N need not be the same size as M .
In contrast to GANs and VAEs, NFs explicitly encode an estimate for the probability density p X .
These networks rely on a coordinate transformation which maps the prior distribution p Z into a target distribution p g with g now being invertible.
This requires M = N but allows for an analytic expression for the probability density induced by g :
[Equation]
In order to match p g and the data probability density p X , one can directly maximize the log-likelihood of the data without resorting to an auxiliary task:
[Equation]
2.2 Topological obstructions
While the bijective nature of NFs allows for an explicit representation of the target probability density estimate, they inevitably suffer from a signiﬁcant drawback.
In order to ﬁnd a mapping g which satisﬁes Eq.
(1) and matches the data probability density, the manifolds deﬁned by the prior distribution p Z and the target distribution p X need to be topologically equivalent.
A common way to describe the topological properties of manifolds are the Betti numbers.
Deﬁnition 1.
The n th Betti number b n of a topological space X is deﬁned as the rank of the n th homology group H n of X [84, 85].
Informally, the n th Betti number denotes the number of n -dimensional holes of a topological space.
In fact, the ﬁrst three Betti numbers do have simple geometrical interpretations:
• b 0 is the number of connected components.
• b 1 is the number of one-dimensional or circular holes.
• b 2 is the number of two-dimensional voids.
For instance, in Fig.
2 we illustrate a torus which has one connected component b 0 = 1 , two one-dimensional holes b 1 = 2 , and a single void enclosed within the surface b 2 = 1 .
It has been proven by Poincaré that these numbers are invariant under homeomorphisms.
This implies the following proposition:
Proposition 1.
Any bijective mapping g ( z ) is a homeomorphism and preserves the topological structure of the input space.
A proof can be found in Refs.
This mathematical statement has a major impact on any generative model relying on bijective mappings, such as NFs.
If the target space has a different topological structure than the latent space, these models cannot properly describe the target space topology as they inherently preserve the latent space topology.
It is indeed possible that these inconsistencies are hidden in some highly-correlated observable which might not be investigated during or after optimization.
Moreover, it has been shown in Ref.
[88] that you can hide and diminish the topological issues if you increase the network complexity.
However, you can achieve better and computationally cheaper results by relaxing the bijetivity constraint itself [88] or by augmenting additional dimensions [87, 89].
Further, the issues arising by non-trivial topologies have also been investigated in the context of autoencoders in Ref.
Consequently, if we want to correctly reproduce the target space with a non-trivial topology, we either need to modify the latent space to match the topological structure, or the bijective model g has to be replaced by a non-topology preserving model such as GANs or VAEs.
While these models do have more ﬂexibility and do not suffer from the same topological obstructions, this freedom usually goes hand in hand with training instabilities and more time-consuming hyperparameter searches.
In order to beneﬁt from both worlds and to ﬁx the topology of the latent space, we propose to use the L A S E R protocol which is described in the following.
Proposed methods
The L A S E R protocol is illustrated in Fig.
The input is a generator g ( z ) that could either be a neural network or a black-box (BB) simulator (e.g. a physics-based Monte Carlo).
As a ﬁrst step, a classiﬁer is trained to distinguish samples drawn from g and samples drawn from the target data probability density.
It is well-known [91, 92] (and forms the basis of the
D CTR protocol [93]) that an ideally trained classiﬁer f using the binary cross entropy (BCE) loss function will have the property f/ (1 − f ) ∝ p ( x | target ) /p
( x | from g ) .
The proportionality constant is p ( target ) /p
( from g ) and is unity when the two datasets have the same number of samples.
Other loss functions such as the mean squared error have the same property.
It is also possible to engineer the loss function to directly learn the likelihood ratio [94, 95].
Throughout this paper, we use BCE.
Each generated sample is assigned a weight w ( x i ) =
f ( x i ) / (1 − f ( x i )) .
These weights are then assigned to the corresponding latent space point g
( z i ) =
x i .
When g is surjective, w is a proper function of z .
However, there are many potential applications where the generator has many internal latent degrees of freedom that are marginalized over and are not part of Z .
In this case, one can use the O MNI F OLD algorithm [75] which is an iterative expectation maximization-style approach that returns weights for the z i that maximizes the likelihood in the data space.
The weights extracted with
O MNI F OLD are a proper function of the latent space.
This is not necessarily the case for weights directly inferred from the classiﬁer, since the mapping g can send the same latent space point to different data space points.
In other words, the marginalized latent dimensions make g a stochastic function.
The ﬁnal step of L A S E R requires to sample from weighted latent space ( z, w ( z )) .
This can be for instance done directly using the following methods:
• Rejection sampling: We can use the weights w ( z ) to perform rejection sampling on the latent space.
This is simple, but ineffective if the weights are broadly distributed.
Thus, we do not employ this method in our experiments.
• Markov chain Monte Carlo (MCMC): The weights w ( z ) induce a new probability distribution
[Equation]
with some normalisation factor Z 0 .
In general, this probability distribution is non-trivial and does not allow to determine its quantile function analytically.
However, if the probability distribution q Z and its derivative ∂q Z /∂z are tractable a MCMC algorithm can be employed.
While Ref.
[72] uses a simple Langevin dynamics algorithm, we suggest to use the more advanced Hamiltonian Monte Carlo (HMC) [96, 97] algorithm as it generates signiﬁcantly better samples.
However, instead of using the above approaches, we suggest for the L A S E R protocol to convert the weighted latent space ( z, w ( z )) into a new unweighted latent space.
For this, we train a second generative model Φ which we call the reﬁner.
The reﬁner maps an auxiliary latent space Y ⊆
R L onto the reﬁned latent space Z (cid:48) ⊆
R N , with underlying probability densities p Y and p Z (cid:48)
=
q Z , respectively.
The reﬁner Φ can be either a GAN or a NF, where the latter requires L = N .
In order to train the model on the weighted latent space ( z, w ( z ))
we accommodate the weights into the loss functions as it was proposed in Refs.
In contrast to the MCMC algorithm, this method does not require to calculate the derivative of the weights ∂w/∂z .
At the end, we generate new samples from the reﬁned generator by ﬁrst sampling from the auxiliary latent space, mapping that to the reﬁned latent space, and then passing that through g .
In some cases, the entire procedure can be iterated, i.e. if g is a neural network, its weights could be updated using the reﬁned latent space.
2 illustrates some of the features of various combinations of generators g and latent space reﬁner networks Φ .
Numerical examples
We consider three different 2-dimensional examples which allows to visualize both the full output and the latent space of the primary generator.
All three sets of training data consist of 480k data
[Table] points. In all models and experimentes we used a batch size of 2000 points and an epoch is deﬁned as iterating through all points in the training set, giving 240 model updates per epoch.
We run all our experiments on our internal GPU cluster, which consists of multiple Nvidia GTX 1080 TI and Nvidia
RTX 2080 Ti GPUs.
Depending on the model the computation times range between 30 mins andhours.
4.1 Implementation details
All models and the training routines are implemented in P Y T ORCH 1.8 [100].
In the following we specify the implementation details for the various networks and the chosen hyperparameters.
Baseline model.
In all examples we use the Real-NVP [101] implementation of a NF as the baseline model.
Our implementation relies partially on the F R EIA library.
In the ﬁrst two examples, the NF consists of 12 coupling blocks, where each block consists of a fully connected multi-layer perceptron (MLP) with 3 hidden layers, 48 units per layer and leaky ReLU activation functions.
In the third more complex example, we increase the number of coupling blocks to 20 and the number of units to 60.
In all cases the model is trained for 100 epochs by minimizing the negative log-likelihood with an additional weight decay of 10 − 5 .
The optimization is performed using Adam [102] with default β parameters and α {+
+} = 10 − 3 .
We further employ a exponential learning rate schedule, i.e. α n = α 0 γ n , with γ = 0 .
999 which decays the learning rate after each epoch.
Post-hoc classiﬁer.
As a post-hoc classiﬁer we employ a MLP with 8 hidden layers, 96 units per layer and leaky ReLU activation functions.
We train it for 50 epochs by minimizing the BCE loss between true data samples and samples from the generator g .
Optimization is again performed using
Adam with default β parameters and α
=
10 − 3 and an exponential learning rate schedule with γ = 0 .
Reﬁner network.
Finally, the reﬁner network is a GAN consisting of a discriminator and a generator, both MLPs with 7 hidden layers, 100 units per layer and Leaky ReLU activation functions.
As a GAN does not require to have the same dimensions in the latent space as in the feature space we choose 4-dimensional auxiliary latent space to simplify the task of the generator.
The standard GAN objective is replaced by the weighted BCE loss introduced in Ref.
[98] to accommodate the weights of the weighted training data.
The GAN is trained for 200 epochs using the Adam optimizer with ( α, β 1 , β 2 ) =
(10 − 4 , 0 . 5 , 0 . 9) and also using an exponential learning rate schedule with γ = 0 .
To compensate an imbalance in the training we update the discriminator four times for each generator update.
Hamiltonian Monte Carlo (HMC).
In order to have full control of all parameters and inputs we implemented our own version of Hamiltonian Monte Carlo in P Y T ORCH , using its automatic differentiation module to compute the necessary gradients to run the algorithm.
For all examples, we initialize a total of 100 Markov Chains running in parallel to reduce computational overhead and solve the equation of motions numerically using the leapfrog algorithm [97] with a step size of ε = 0 .
004 and 50 leapfrog steps.
To avoid artifacts originating from the random initialization of the chains, we start with a burn-in phase and discard the ﬁrst 3000 points from each chain.
4.2 Probability distance measures
In order to quantify the performance of the base model and its reﬁned counterparts, we implemented two measures of distance:
Jensen–Shannon divergence (JSD).
A simple but commonly used distance measure between two probability distributions p ( x ) and q ( x ) is the Jensen–Shannon divergence
[Equation]
In contrast to the Kullback–Leibler (KL) divergence it is symmetric and always returns a ﬁnite value.
Earth mover distance (EMD).
Another common distance measure between two probability distributions p ( x ) and q ( x ) is the earth mover’s distance (EMD).
The EMD is usually formulated as an optimal transport problem.
Lets assume the two distributions are represented as clusters
P = { ( x i , p ( x i ) } Mi =1 and Q = { ( y j , q ( y j ) } M (cid:48)
j =1 .
Then, the EMD between two probabilities is the minimum cost required to rearrange one probability distribution p into the other q by discrete probability ﬂows f ij from one cluster point x i to the other y j
[Equation]
M (cid:48) (cid:88) j = q ( y j )   , (5) where the optimal ﬂow f ij is found by solving this optimization problem.
For an efﬁcient implementation we used the POT library [103].
Numerical approximation.
In order to calculate these quantities numerically we generate 2M data points for all models and examples and approximate the probability densities from a 2-dimensional histogram.
As ground distance measure between clusters P and Q for the calculation of the EMD, we consider the Euclidean distance between the center points of each bin.
In order to estimate the error by approximating the probabilities, we calculated the JSD and EMD score on equally large but independent test samples drawn from the truth distribution.
The extracted error estimations are indicated in parenthesis in Table 3.
4.3 Experimental results
In the following we will simply refer to the L A S E R protocol using the Hamiltonian Monte Carlo as
HMC method, and we will refer to the L A S E R protocol using a reﬁner network as L A S E R method.
In the ﬁrst example, illustrated in Fig.
3, we considered a multi-modal distribution which can be topologically described by its Betti numbers b 0 = 4 and b n ≥ 1 = 0 .
In contrast, our non-reﬁned latent space is described by a 2-dimensional normal distribution with b 0 = 1 and b n> 0 = 0 .
As both spaces are topologically different our baseline model fails to correctly reproduce the truth distribution.
It is obvious that the baseline model cannot properly separate the individual modes and always leaves a connection between them.
This is in agreement with Proposition 1 as our baseline model is a NF and thus inevitably topology preserving.
After training a post-hoc classiﬁer the reweighted feature space of the D CTR method clearly resolved the topological obstructions.
Pulling back these weights onto the orginal latent space induces a weighted latent space which is now also clearly separated into four pieces and thus shows the same topological structure as the data distribution.
This weighted latent space is correctly reproduced in the reﬁned latent spaces of the HMC and L A S E R method.
Acting with our generator g on these reﬁned latent spaces shows a highly improved feature space distribution.
In Fig.
4, we considered a double donut distribution with Betti numbers b 0 = 1 , b 1 = 2 and b n ≥ 2 =as a second example.
This example does not have disconnected parts but contains two circular holes.
Again the baseline model has troubles to properly close the loops and reproduce the holes as these topological features are not present in the non-reﬁned latent space. After reweighting the latent space [Figure]
Dctr [Figure] with the D CTR weights, these topological features emerge as holes in the latent space.
Both reﬁned latent spaces show the same topological structure as the reweighted latent space, resulting in an improved feature space distribution for both the HMC and the L A S E R method.
[Figure]
Finally, in the last example we combine the topological features from the previous examples to make our baseline model fail even more.
In detail, we considered a set of 3 displaced rings which can be topologically described by the Betti numbers b 0 = 3 , b
1 = 3 and b n ≥ 1 = 0 , as illustrated in Fig.
As with previous examples, the baseline model utterly fails to reproduce the true data distribution.
As the topological structure of the data distribution is more complex the reweighted latent space in the D CTR method is considerably more involved then the weighted latent spaces of the previous examples.
Owing to this more complex structure of the weighted latent space our reﬁner model Φ has a much harder job to properly reproduce these topological features in the reﬁned latent space of the
L A S E R method.
While being far from optimal the feature space distribution of the L A S E R reﬁned generator is notably improved in comparison to the baseline model. In this particular complicated example the output of the HMC method seems to show slightly better agreement with the truth distribution than the L A S E R method. However, these results can be improved by spending more time on a detailed network hyperparameter optimization. In order to provide a quantitative comparison of performances, we calculated the EMD and JSC scores of our baseline model and its reﬁned counterparts for all examples.
These scores are summarized in Table 3 and show that the HMC method yields the best performance in most scenarios for the unweighted output.
However, the L A S E R method is equally good or better for the the most complex example, showing that the HMC method has troubles to sample from the highly tangled latent space.
[Figure]
Truth
Baseline
HMC
LaSeR
Dctr
[Figure] [Table]
Moreover
[Table]
the HMC method cannot be used at all if the derivatives of the weights are not tractable, in which case the L A S E R method is still applicaple.
The overall best performance is achieved by the D CTR method which however corresponds to a weighted feature space.
As these weights are also used to obtain the reﬁned latent space for the HMC and L A S E R method the score of the D CTR yield as a benchmark and provides lower limit.
Indeed, we can ask ourselves why the EMD of both unweighted methods are smaller than the EMD of the D CTR method for the double donut example.
However, owing to the the uncertainty on the numerical estimation of the scores, which are indicated in parenthesis, all scores a statistically consistent.
Conclusions and outlook
We have presented the L A S E R protocol, which provides a post-hoc method to improve the performance of any generative model.
In particular, it has shown remarkable results on several numerical examples and improved the EMD (JSD) score of the baseline model by 25-61% (62-67%).
It further managed to ﬁx the topological obstructions encountered by the base model.
While we have only explicitly demonstrated L A S E R for reﬁning normalizing ﬂows as the primary generator, the protocol is also valid for other deep generative models as well as other generative models that need not be differentiable or surjective.
While L A S E R is a promising approach for generative model reﬁnement, it also has some limitations.
First of all, there is some dependence on the efﬁcacy of the primary generator.
For example, if the primary generator is a constant, then no amount of reﬁning can recover the data probability density.
Furthermore, the reﬁnement comes at a computational cost with an additional step on top what is required to construct the original generative model.
This can be particularly important when the new latent space is sourced from a much bigger space than the original one.
Finally, even though L A S E R produces unweighted examples, the statistical precision of the resulting dataset may be limited by the size of the training dataset [42].
We leave further investigations of the statistical power of reﬁned samples and the interplay with hyperparameter choices for future research.
Acknowledgments and Disclosure of Funding
BPN was supported by the Department of Energy, Ofﬁce of Science under contract number DEAC02-05CH11231.
RW acknowledges support by HeiKA and by FRS-FNRS
(Belgian National
Scientiﬁc Research Fund) IISN projects 4.4503.16.
MB is supported by the International Max Planck
School Precision Tests of Fundamental Symmetries.
We thank Tilman Plehn, Uroš Seljack and Jesse
Thaler for their helpful feedback on the manuscript