E MBEDDED - MODEL FLOWS : C OMBINING THE INDUCTIVE BIASES OF MODEL - FREE DEEP LEARNING AND EXPLICIT PROBABILISTIC MODELING
A BSTRACT
Normalizing flows have shown great success as general-purpose density estimators.
However, many real world applications require the use of domain-specific knowledge, which normalizing flows cannot readily incorporate.
We propose embedded-model flows (EMF) which alternate general-purpose transformations with structured layers that embed domain-specific inductive biases.
These layers are automatically constructed by converting user-specified differentiable probabilistic models into equivalent bijective transformations.
We also introduce gated structured layers , which allow bypassing the parts of the models that fail to capture the statistics of the data.
We demonstrate that EMFs can be used to induce desirable properties such as multimodality and continuity.
Furthermore, we show that EMFs enable a high performance form of variational inference where the structure of the prior model is embedded in the variational architecture.
In our experiments, we show that this approach outperforms a large number of alternative methods in common structured inference problems
I NTRODUCTION
Normalizing flows have emerged in recent years as a strong framework for high-dimensional density estimation, a core task in modern machine learning (Papamakarios et al., 2021; Kobyzev et al., 2020; Kingma & Dhariwal, 2018).
As with other deep generative architectures such as GANs and
VAEs, normalizing flows are commonly designed with the flexibility to model a very general class of distributions.
This works well for large naturalistic datasets, such as those common in computer vision.
However, general-purpose flows are less appealing when datasets are small or when we possess strong a priori knowledge regarding the relationships between variables.
On the other hand, differentiable (or ‘deep’) probabilistic programming is a powerful and general framework for explicit stochastic modeling that can express strong assumptions about the datagenerating process (Tran et al., 2017; 2016;
Bingham et al., 2019; Dillon et al., 2017; Piponi et al., 2020; Kucukelbir et al., 2017; Ambrogioni et al., 2021a).
For example, the user may know that the data represents a process evolving over time or in space, or is subject to known physical laws, or that there is an unobserved common cause or other shared factor behind a set of observations.
Such structured prior information can enable the user to form effective inferences using much less data than would otherwise be required.
For example, in the field of astrophysics, differentiable probabilistic programs have been successfully used to study strong gravitational lensing (Chianese et al., 2020) and dark matter substructures (Coogan et al., 2020; Varma et al., 2020).
All of these applications entail the translation of sophisticated physical theories into highly parameterized and structured probabilistic programs.
However, it remains the case that “all models are wrong” (Box, 1976), and it is rare that a modeler can be fully confident in the strong assumptions (e.g., Gaussianity, statistical independence) expressed by typical parametric probability models.
In this paper we bridge these two worlds, providing an automatic technique to convert differentiable probabilistic programs into equivalent normalizing flow layers with domain-specific inductive biases.
We call these architectures embedded-model flows (EMF).
EMFs can be use to integrate the structure of “wrong but useful” models within generic deep learning architectures capable of correcting the deviation of the model from the data.
Furthermore, in the context of Bayesian inference, we show that the EMF architecture can be used to define a powerful class of approximate posterior distributions that embed the structure of the prior.
P RELIMINARIES
Normalizing flows define a random variable as an invertible differentiable transformation y = f ϕ ( x (0) ) of a base variable x (0) ∼ p {+
+} ( x (0) ) .
In these expressions, ϕ are parameters that control the form of the transformation and consequently the distribution of y .
The transformed (log-)density can then be given explicitly using the change of variables formula from probability theory,
[Equation]
where J ϕ ( y ) is the Jacobian matrix of f ϕ .
The base distribution p 0 is often taken to be standard normal ( x (0) ∼ N ( 0 , I ) ).
In typical applications, the functional form of f ϕ is specified as the composition of a series of simpler invertible transformations: [Equation] We denote the resulting sequence of transformed variables as x (0) , x (1) , . . . , x ( m − 1) , y .
Each of these transformations (i.e. flow layers ) are designed to have tractable formulas for inversion and log-Jacobian.
The parameters ϕ are usually trained by stochastic gradient descent to maximize the log-likelihood (given observed data) or evidence lower bound (in the variational inference setting).
2. P ROBABILISTIC PROGRAMS AND DIRECTED GRAPHICAL MODELS
Probabilistic programming allows a joint probability distribution over a set of random variables to be specified intuitively as a program that includes random sampling steps (van de Meent et al., 2018).
Computations on this distribution, including the joint log-density function, structured inference algorithms, and the bijective transformations described in this paper, may then be derived automatically as program transformations.
Recently frameworks have emerged for ‘deep’ probabilistic programming, which exploit the automatic differentiation and hardware acceleration provided by deep learning frameworks; examples include Pyro (Bingham et al., 2019), PyMC (Salvatier et al., 2016), Edward (Tran et al., 2017), and TensorFlow Probability (Dillon et al., 2017;
Piponi et al., 2020).
These allow for straightforward implementation of gradient-based inference and parameter learning, and enabling methods such as ours that integrate probabilistic programs with standard ‘deep learning’ components such as flows.
Probabilistic programs that contain only deterministic control flow may equivalently be viewed as directed graphical models, in which each random variable x
j is a graph node connected by incoming edges to some subset of the previous variables x <j .
This leads to the following general expression for the joint density of a set of vector-valued variables x 1: n = { x j } nj =1 :
[Equation]
Here θ j ( x <j ; ϕ ) denotes the ‘link function’ associated with each node that takes as input the model parameters ϕ and the values of all parent variables (excepting θ 0 , where there are no parents), and outputs the parameters of its respective density function.
As a practical matter, flow architectures generally assume layers of fixed size, so the experiments in this paper focus on programs of fixed structure and we use graphical-model notation for simplicity.
However, the methods we describe do not require access to an explicit graph structure and could be straightforwardly applied to programs with stochastic control flow.
C ONVERTING PROBABILISTIC PROGRAMS INTO NORMALIZING FLOWS
In this section, we will show that a large class of differentiable probabilistic programs can be converted into equivalent normalizing flow layers.
More precisely, we will construct an invertible function F ϕ that maps spherical normal random inputs to the joint distribution of the probabilistic program:
[Equation]
This is a special case of the so called Rosenblatt transform (Rosenblatt, 1952).
We interpret this transform as a flow layer and we extend it by including trainable gate variables that can be used to ”disconnect” a set of variables from its parents so to correct model miss-specification.
In our applications, the input to the bijective transformation may be arbitrary real-valued vectors, for example, the output of a previous flow layer, which likely will not be normally distributed.
In such cases the transformed values will not be samples from the original probabilistic program, but we expect the transformation to capture some of its properties, such as multimodality and the dependence structure of its random variables.
3. I NDIVIDUAL RANDOM VARIABLES AS BIJECTIVE TRANSFORMATIONS
We first consider the case of a univariate random variable x
(1) ∼ p ϕ
( x (1) ) , which we can express as a continuous bijective transformation f ϕ of a standard normal variable x (0) via a slight extension of inverse transform sampling.
The probability integral transform theorem states that the cumulative distribution function (CDF) of p ϕ , C ϕ ( x ) = (cid:82) x a p ϕ ( y )
d y , with p ϕ defined over a generic interval ( a, b ) and a < x < b , is uniformly distributed, u = C ϕ ( x ) ∼ U (0 , 1) (Angus, 1994).
If the CDF is invertible, i.e., strictly increasing within its domain, it follows that we can sample x = C − 1 ϕ ( u ) ∼ p ϕ given only a uniform random variable u ∼ U (0 , 1) .
We obtain u through a second application of the probability integral transform theorem to the standard normal CDF Φ 0 , 1 .
Our univariate ‘flow’ is therefore given by
[Equation]
Closed-form expressions for f ϕ may be available in some cases (see below).
In the general univariate setting where the inverse distribution function C − 1 ϕ ( u ) is not available in closed form, it may be evaluated efficiently using numerical root-finding methods such as the method of bisection (see Appendix A.6), with derivatives obtained via the inverse function theorem and the implicit reparameterization trick (Figurnov et al., 2018).
However, this is not needed for maximum likelihood training of the parameters ϕ , since eq.
(1) only involves f − 1 ϕ and its derivatives.
Example: Gaussian variables .
The inverse CDF of a Gaussian variable x (1) ∼ N (cid:0) µ, σ 2 (cid:1) can be expressed as a scale-shift transformation of the inverse CDF of a standard normal variable:
Φ − 1 µ,σ ( u ) = µ + σ Φ − 1 0 , 1 ( u ) .
We can obtain the forward flow transformation by composing this function with the CDF of a standard Gaussian: f µ,σ ( x (0) ) =
Φ − 1 µ,σ (Φ 0 , 1 ( x (0) ))
= µ + σx (0) .
This is the famous reparameterization formula used to train VAEs and other variational methods (Kingma & Welling, 2013).
The inverse function is f − 1 µ,σ
( x (1) ) = ( x (1) − µ ) /σ while the log-Jacobian is equal to log σ .
It is not generally straightforward to evaluate or invert the CDF of a multivariate distribution.
However, the most popular multivariate distributions, such as the multivariate Gaussian, can be expressed as transformations of univariate standard normal variables, and so may be treated using the methods of this section.
More generally, the next section presents a simple approach to construct structured bijections (i.e., flows) from multivariate distributions expressed as probabilistic programs.
3. F ROM DIFFERENTIABLE PROBABILISTIC PROGRAMS TO STRUCTURED LAYERS
A probabilistic program samples a set of variables sequentially, where (as in eq.
the distribution of each variable is conditioned on the previously-sampled parent variables.
Using the results of Section 3.1, we can represent this conditional sample as an invertible transformation of a standard normal input.
Feeding the result back into the model defines a sequence of local transformations, which together form a multivariate flow layer (algorithm 1).
More specifically, as with autoregressive flows, the forward transformation of the j -th input requires us to have computed the preceding outputs x ( i +1) < j .
This may be expressed recursively as follows:
[Equation]
where f ( i ) j, ϕ ( · ; x ( i +1) <
j ) denotes the invertible transformation associated to the conditional distribution of the j -th variable given the values of its parents.
The recursive expression terminates since the parenting graph has an acyclic-directed structure.
In the reverse direction, the variables x ( i +1) 1: n are all given, so all transformations may be evaluated in parallel:
[Equation]
This is a generalization of the inversion formula used in autoregressive flows (Kingma et al., 2016;
Papamakarios et al., 2017).
The Jacobian of the transformation is block triangular (up to a permutation of the variables) and, consequently, the log determinant is just the sum of the log determinants of the block diagonal entries:
[Equation]
where J − 1 j,ϕ ( x ( i +1) j ; x ( i +1) <j ) is the Jacobian of the local inverse transformation [Equation]
We denote this multivariate invertible transformation as a structured layer . If the input x ( i ) 1: n follows a standard Gaussian distribution, then the transformed variable follows the distribution determined by the input probabilistic program. This is not going to be the case when the layer is placed in the middle of a larger flow architecture. However, it will still preserve the structure of the probabilistic program with its potentially complex chain of conditional dependencies.
Gated structured layers Most user-specified models offer a simplified and imperfect description of the data. In order to allow the flow to skip the problematic parts of the model, we consider gated layers. In a gated layer, we take each local bijective transformation g j, ϕ to be a convex combination of the original transformation f j, ϕ and the identity mapping: [Equation]
[Equation] structure; { ( i ) k } = procedure G ATED F ORWARD S AMPLER ( { x ( i ) k } nk =1 ) State log | det J | ←for j in Range( n ) do { p l } m j l =1 ← x ( i [-=
-]) j .parents idx ▷ { p l } m
j l =1 : indices of parents of j -th variable.
▷ m j : number of parents of j -th variable.
if { p l } m j l =1 is not empty then ˜ x
( i +1)
j ← f ( i ) j, ϕ ( x ( i ) j ; { x ( i +1) p l } m j l =1 ) x
( i +1)
j ← λ ( i ) j ˜ x
( i +1) j + (1 − λ ( i ) j ) x
( i )
j log | det J | ← log | det J | + log (cid:16) det | J j,ϕ ( { x ( i +1) k } nk =1 ) + (1 − λ ( i ) j ) I | (cid:17) else
˜ x
( i +1)
j ← f ( i ) j, ϕ
( x ( i ) j ; ϕ j ) ▷ ϕ j : parameters of the j -th root distribution.
x ( i +1)
j ← λ ( i ) j ˜ x
( i +1)
j + (1 − λ ( i ) j ) x
( i )
j log | det J | ← log | det J | + log (cid:16) det | J j, ϕ j 0 ( x ( i +1) j )
+ (1 − λ ( i ) j ) I | (cid:17) return { x ( i +1) k } nk =1 , log | det J |
This gating technique is conceptually similar to the method used in highway networks (Srivastava et al., 2015), recurrent highway networks (Zilly et al., 2017) and, more recently, highway flows (Ambrogioni et al., 2021b).
Here the gates λ ( i ) j ∈ (0 , 1) are capable of decoupling each node from its parents, bypassing part of the original structure of the program.
As before, the local inverses g − 1 ϕ ( x ( i +1) j ; x ( i +1) < j , λ ( i ) j ) can in general be computed by numeric root search, but the Gaussian case admits a closed form in both directions:
[Equation]
Analogues of eq.
(6), eq. (7), and eq.
(8) then define the gated joint transformation, its inverse and log Jacobian determinant, respectively.
Note that the gating affects only the diagonal of the Jacobian, so the overall block-triangular structure is unchanged.
3. A UTOMATION
All the techniques described in this section can be fully automated in modern deep probabilistic programming frameworks such as TensorFlow Probability.
A user-specified probabilistic program in the form given by Eq.
3 is recursively converted into a bijective transformation as shown in the pseudo-code in Fig.
In TensorFlow Probability, the bijector abstraction uses pre-implemented analytic inverse and Jacobian formulas and can be adapted to revert to root finding methods otherwise.
3. E MBEDDED - MODEL FLOWS
A EMF is a normalizing flow architecture that contains one or more (gated) structured layers.
EMFs can be used to combine the inductive biases of explicit models with those of regular normalizing flow architectures.
Several probabilistic programs can be embedded in a EMF both in parallel or sequentially.
In the parallel case, a single flow layer is comprised by two or more (gated) structured layers acting on non-overlapping subsets of variables.
The use of several programs embedded in a single architecture allows the model to select the most appropriate structure during training by changing the gating parameters and the downstream/upstream transformations.
This allows the user to start with a ”wrong but useful model” and then learn the deviation from the distribution of the collected data using flexible generic normalizing flow transformations.
A visualization of a hypothetical EMF architecture is given in Fig.
R ELATED WORK
Embedded-model flows may be seen as a generalization of autoregressive flows (Kingma et al., 2016; Papamakarios et al., 2017) in the special case where the probabilistic program consists of a Gaussian autoregression parameterized by deep link networks, although generic autoregressive architectures often permute the variable order after each layer so that the overall architecture does not reflect any specific graphical structure.
Graphical normalizing flows (Wehenkel & Louppe, 2021) generalize autoregressive structure to arbitrary DAG structures; however, this approach only constrains the conditional independence structure and does not embed the forward pass of a user-designed probabilistic program.
In the last few years, the use of tailored model-based normalizing flow architectures has gained substantial popularity.
Perhaps the most successful example comes from physics, where gaugeequivariant normalizing flows are used to sample from lattice field theories in a way that guarantees preservation of the intricate symmetries of fundamental fields (Albergo et al., 2019; Kanwar et al., 2020; Albergo et al., 2021;
Nicoli et al., 2021).
Similarly, structured flows have been used to model complex molecules (Satorras et al., 2021), turbulence in fluids (Bezgin & Adams, 2021), classical dynamical systems (Rezende et al., 2019; Li et al., 2020) and multivariate timeseries (Rasul et al., 2020).
These approaches are tailored to specific applications and their design requires substantial machine learning and domain expertise.
Normalizing flows have been applied to variational inference since their introduction by Rezende &
Mohamed (2015).
Recognizing that structured posteriors can be challenging for generic architectures, a number of recent approaches attempt to exploit the model structure.
Structured conditional continuous normalizing flows (Weilbach et al., 2020) use minimally faithful inversion to constrain the posterior conditional independence structure, rather than embedding the forward pass of the prior model.
Our architecture can be seen as a flexible extension of previous variational models that also embed the prior structure, such as stochastic structured variational inference (Hoffman & Blei, 2015), automatic structured variational inference (Ambrogioni et al., 2021a) and cascading flows (Ambrogioni et al., 2021b).
It is also closely related to the non-centering transformations applied by
Gorinova et al.
(2020) to automatically reparameterize a probabilistic program, although that work focused on improving inference geometry and did not directly exploit the bijective aspect of these reparameterizations.
A PPLICATIONS
Model-embedding allows users of all levels of expertise to insert domain knowledge into normalizing flow architectures.
We will illustrate this with some examples, starting from simple uncoupled multi-modal models and moving on to more sophisticated structured models for time series.
Finally, we will discuss the application of EMF to automatic structured variational inference, where we use EMF architectures that embed the prior distribution.
Details about the datasets used, the models’ hyperparameters and the training procedures can be found in Appendix A.1, A.3 and A.4 respectively.
Additional experiments with an hierarchical structure can be found in appendix A.7.3.
Number of parameters and sampling time for all the models are reported in appendix A.5.
All the code used in the experiments is available at https://github.com/gisilvs/EmbeddedModelFlows .
5. M ULTIMODALITY
Multimodal target densities are notoriously hard to learn for traditional normalizing flow architectures (Cornish et al., 2020).
In contrast, it is straightforward to model multimodal distributions using explicit mixture models; the simplest example is perhaps a model where each variable follows an independent mixture of Gaussian distributions:
[Equation]
The CDF of a mixture of distributions is given by the mixture of the individual CDFs.
Therefore, the inverse transformation associated with a mixture of Gaussians model is
[Equation]
In itself, the probabilistic program in Eq.
12 can only model a very specific family of multivariate mixture distributions with Gaussian components organized in an axis-aligned grid.
However, the model becomes much more flexible when used as a structured layer in an EMF architecture since the generic upstream layers can induce statistical coupling between the input variables, thereby loosening the independence structure of the original model.
The resulting approach is conceptually similar to architectures based on monotonic transformations such as neural spline flows (Durkan et al., 2019a;b), which are arguably more efficient as they have an analytic inversion formula.
However, this application shows that a specific inductive bias such as multimodality can be effectively incorporated into a flow architecture by users with only a basic understanding of probabilistic modeling.
We apply an EMF with an independent mixture of Gaussians to two 2D toy problems common in the flow literature, namely the “eight Gaussians” and the “checkerboard”, and to the MNIST dataset.
As baselines, we use a Masked Autoregressive Flow (MAF) model with two autoregressive layers and a standard Gaussian as a base density, a “large” version of MAF (MAF-L) with three autoregressive layers, which has much more flexibility compared to the other models in terms of trainable parameters, and a six steps Neural Splines Flows (NFS, with Rational Quadratic splines) with coupling layers as proposed in (Durkan et al., 2019b).
We then combine the EMF with the MAF baseline by adding an
EMF on top of the two autoregressive layers (EMF-T, for top) and in between the two layers, after the latent variables permutation (EMF-M, for middle).
We also combine NSF with the structured layer, both with the EMF-T and EMF-M settings.
Each mixture has 100 components, with trainable weights, means and standard deviations.
The results are reported in table 1, while densities are shown in appendix A.7.1.
The use of the multimodal structured layer with MAF always outperforms the autoregressive models.
The improvement is not solely due to the increased parameterization as both methods outperform the more complex baseline MAF-L. NSF achieves superior results, as it is an architecture designed to model multimodality.
The combination of the structured layer with NSF achieves the overall best performances.
However, this comes at the cost of additional sampling time, caused by the numeric inversion of the mixture of Gaussians transformation.
[Table]
5. T IMESERIES MODELING
Timeseries data is usually characterized by strong local correlations that can often be described with terms such as continuity and smoothness.
In the discrete models considered here, these properties can be obtained by discretizing stochastic differential equations.
One of the simplest timeseries model with continuous path is the Wiener process ˙ x ( t ) =
w ( t ) , with w ( t ) being a white noise input with variance σ 2 .
By Euler–Maruyama discretization with ∆ t = 1 , we obtain the following autoregressive model
[Equation]
(13) which gives us the bijective function f t +1 ,σ ( x t +1 ; x t ) =
x t + σx t +1 .
We use four stochastic differential equation models as time series datasets, discretized with the Euler-Maruyama method: the models are Brownian motion (BR), Ornstein-Uhlenbeck (OU) process, Lorenz system (LZ), and
Van der Pol oscillator (VDP).
We use an EMF with continuity structure.
The baselines are MAF,
MAF-L and NSF models like in the previous section, and an MAF in which the base distribution follows the continuity structure (B-MAF), as a discrete-time version of the model proposed in (Deng et al., 2020).
Note that no variables permutation is used for B-MAF.
We then use GEMF-T and a combination of NSF with the structured layer (NSF-GEMF-T), with continuity structure.
The results are reported in Table 2.
The baselines combined with continuity structure greatly outperforms the other models (even the more complex MAF-L) for all of the datasets.
Results obtained with the smoothness structure are reported in appendix A.7.4.
[Table]
5. V ARIATIONAL INFERENCE
Normalizing flows are often used as surrogate posterior distributions for variational inference in
Bayesian models, where the target posterior may contain dependencies arising both from the prior structure and from conditioning on observed data via collider (or “explaining away”) phenomena.
Generic flow architectures must learn the prior structure from scratch, which wastes parameters and computation and may lead the optimization into suboptimal local minima.
Embedding the prior model into the flow itself as a gated structured layer avoids this difficulty and should allow the remaining layers to focus their capacity on modeling the posterior.
We deploy a gated EMF architecture with a two-layers of inverse autoregressive flow (IAF) and a final structured layer embedding the prior model.
This is visualized in 2 for inference in a Lorentz dynamical system.
In addition, we train plain IAF, automatic differentiation Mean Field (MF) and
Multivariate Normal (MVN) surrogate posteriors from Kucukelbir et al.
(2017) and ASVI from
Ambrogioni et al.
(2021a).
We experiment on both time series and hierarchical models.
The time series models are Brownian motion (BR), Lorenz system (LZ) and Van der Pol oscillator (VDP), with
Bernoulli (classification) emissions.
Furthermore, we either observe all the emissions (smoothing setting) or we omit a central time window (bridge setting).
As hierarchical models we use the Eight
Schools (ES) dataset and a Gaussian binary tree model with different depths and link functions, in which only the last node is observed (Ambrogioni et al., 2021b).
The results are shown in Table 3.
Results on additional experiments can be found in Appendix A.7.5.
GEMF-T outperforms the baselines on all the problems (on-par with some baselines on the most trivial problems, see Appendix A.7.5), showing the benefits of embedding the prior program for highly structured problems, and the capability of such architecture to model collider dependencies.
[Table]
D ISCUSSION
We introduced EMF, a technique to combine domain-specific inductive biases with normalizing flow architectures.
Our method automatically converts differentiable probabilistic programs into bijective transformations, allowing users to easily embed their knowledge into their models.
We showed how, by choosing appropriate inductive biases, EMF can improve over generic normalizing flows on a range of different domains, with only a negligible increase in complexity, and we achieved high performance on a series of structured variational inference problems.
EMF is a powerful generic tool which can find several applications both in probabilistic inference and more applied domains.
Limitations: I) While we can map a very large family of probabilistic models to normalizing flow architectures, the approach is still limited by the topological constraints of bijective differentiable transformations (Cornish et al., 2020).
This forces the support of the program to be homeomorphic to the range, making impossible to map models with non-connected support sets into normalizing flows.
II) While the root finding methods used for inverting univariate densities are reliable and relatively fast, they do have higher computational demands than closed form inversion formulas.
III) While it is possible to model discrete variables using Eq.
6, this does not result in an invertible and differentiable transformation that can be trained using the standard normalizing flow approach.
IV)
The probabilistic program must be executed during sampling.
This can result in significant sampling slow-downs when using complex autoregressive programs.
Note that, except in the case of variational inference, this does not lead to training slow-downs as the inverse transform is always trivially parallelizable.
Future work: Embedded-model flows are a generic tool that can find application in several domains in which there is structure in the data.
In theory, any number of structured layers can be added to EMF architectures, and combining different structures in the same model may help model complex dependencies in the data.
Further empirical work is needed to asses the capacity of trained EMFs to select the correct (or the least wrong) model.
While we only considered the use of structured layers in normalizing flows, they may also be used in any kind of deep net in order to embed domain knowledge in tasks such as classification, regression and reinforcement learning A CKNOWLEDGEMENTS
OnePlanet Research Center aknowledges the support of the Province of Gelderland.
A A PPENDIX A. D ATASETS
Here we provide a detailed description of the datasets used in the experiments.
A.1. 2 D TOY PROBLEMS
The implementation of the 2D toy datasets is taken from the publicly available implementation of
FFJORD (Grathwohl et al., 2019).
A.1. MNIST
We use as training set 50000 data points, while the remaining 10000 are used as validation set.
The images are dequantized following the same procedure as in (Papamakarios et al., 2017).
A.1. S TOCHASTIC DIFFERENTIAL EQUATIONS
For the generative time series experiments we generate data using four stochastic differential equations (SDEs) models, discretized with the Euler-Maruyama method.
The SDEs are: Geometric Brownian motion, Ornstein-Uhlenbeck process, Lorenz system and Van der Pol oscillator.
In the following, W t corresponds to a Wiener process, and the series length T = 30 with time step s = 1 unless specified.
Geometric Brownian motion : the dynamics evolve as:
[Equation]
where µ and σ are constants.
In practice, we generate sequences as a simple Brownian motion, and then take the exponential of the whole sequence.
The Brownian motion sequence is generated as:
[Equation]
where the mean µ = 0 and the standard deviation σ = 0 . 1 .
Ornstein-Uhlenbeck process : This process has the following dynamics:
[Equation]
with θ > 0 and σ > 0 .
The sequences are generated as:
[Equation]
with µ 0 = 0 , σ 0 = 5 , σ = 0 .
5 and θ = 0 . 8 .
Lorenz system : the Lorenz dynamics evolve as follows:
[Equation]
where ϕ = 10 , ρ = 28 and β =
In the discrete case we have:
[Equation]
with standard deviation σ = 0 .
1 and step size s = 0 .
Van der Pol oscillator : The system evolves as:
[Equation]
In the discrete case, with T = 120 and s = 0 .
05 , we have:
[Equation]
Here we use µ = 1 and σ = 0 .
A.1. S MOOTHING AND B RIDGE DATA
The data for the variational inference experiments is generated with either a Brownian motion, a
Lorenz system or the Van der Pol oscillator.
The first follows the same dynamics of the geometric
Brownia motion described above, but without taking the exponential, while the others are the same as the Lorenz system and Van der Pol oscillator described above.
The true time series remains unobserved, and we only have access to noisy emissions.
In the regression case, the emissions follow a Gaussian distribution:
[Equation]
where σ e is the emission noise, and is σ e = 0 .
15 for the Brownian motion, σ e = 1 for the Lorenz system, and σ e = 0 .
5 for the Van der Pol oscillator.
In the classification case, the emissions follow a
Bernoulli distribution:
[Equation]
with k a gain parameter, k = 5 in the Brownian Motion, k = 2 in the Lorenz system and k = 1 for the Van der Pol oscillator.
In the smoothing problem, we observe emission for all the time steps, while in the bridge problem we observe emissions only in the first and last 10 time points (first and lastfor the Van der Pol oscillator).
For the Lorenz system amd Van der Pol oscillator, the emissions are observed only for the x variables in both smoothing and bridge.
A.1.
[Equation] [Table]
IRIS dataset : We use the IRIS dataset from the UCI Machine learning Repository.
The dataset is composed by three classes of iris plant, each described by four attributes: sepal length, sepal width, petal length and petal width.
There are 50 datapoints for each class.
In the experiments, we use samples of n = 10 datapoints.
Digits dataset : We use the Digits dataset from the UCI Machine learning Repository.
The dataset is composed by 10 classes 8 × 8 images of handwritten digits, from 0 to 9.
In the experiments, we use samples of n = 20 datapoints.
The images are dequantized following the same procedure as in (Papamakarios et al., 2017).
Eight Schools : The Eight Schools (ES) model considers the effectiveness of coaching programs on a standardized test score conducted in parallel at eight schools.
It is specified as follows:
[Equation]
where µ represents the prior average treatment effect, τ controls how much variance there is between schools, i = 1 , . .
, 8 is the school index, and y i and σ i are observed.
Gaussian Binary tree : The Gaussian Binary tree is a reverse tree of D layers, in which the variables at a given layer d , x dj , are sampled from a Gaussian distribution where the mean is function of two parent variables π d − 1 j, 1 , π d − 1 j, 2 at the previous layer:
[Equation]
All the variables at the 0-th layers are sampled from a Gaussian with mean 0, and variance σ 2 = 1 .
The last node of the tree is observed, and the inference problem consists in computing the posterior of the nodes at the previous layers.
We use a linear coupling function link ( x, y ) =
x − y and a nonlinear one link ( x, y ) =
tanh( x ) − tanh( y ) , with trees of depth D
= 4 and D = 8 .
A. S TRUCTURES INDUCED WITH EMF
In this section we describe in detail the structures embedded with EMF and the bijective transformations they correspond to.
In the following, we denote as ¯ x the output vector from the normalizing flow, and ¯ x ′ the vector after the transformation induced by EMF.
When the probabilistic program embedded by EMF has only Gaussian distributions N ( µ, σ 2 ) , the bijective transformation is in the form:
[Equation]
while in the gated case, with λ ∈ (0 , 1) the transformation becomes:
[Equation]
In both time series and Gaussian hierarchical generative models, the embedded structures have only
Gaussian distributions, so all the gated and non-gated bijective transformation will have the form described above.
What changes in the structure is the way we define the means and variances of such structures.
Continuity : for the first variable of the sequence, we have µ = 0 , σ = σ s , while for the remaining
T − 1 variables x t ∀ t ∈ [1 , . . .
, T − 1] , we have µ = x ′ t − 1 and σ = σ s .
σ s is a trainable variable transformed with the softplus function, shared among all the latent variables.
Smoothness : the first two variables of the sequence have µ = 0 and σ = 1 .
The remaining variables have µ = 2 x ′ t − x ′ t − 1 and σ = σ s , with σ s a trainable variable like the case above.
Gaussian hierarchical model : for the first variable of the sequence, corresponding to the mean of the remaining variables, we have µ = 0 and σ = 1 , while for the remaining variables, we set µ = x ′and σ = 1 .
A. F LOW M ODELS
In all the experiments, we use MAF for generative models and IAF for variational inference, both with two autoregressive layers (three for MAF-L) and a standard Gaussian as a base density.
Each autoregressive layer has two hidden layers with 512 units each.
We use ReLU nonlinearity for all the models but for IAF, MAF and MAF-L without EMF or GEMF, for which we empirically found Tanh activation to be more stable.
The features permutation always happen between the two autoregressive layers.
In the case of EMF-M and GEMF-M, the permutation happens before the structured layer.
We found empirically that training the gates for GEMF is slow as the variables lie on an unbounded space before the Sigmoid activation.
A possible solution is scaling the variables (by 100 in our experiments) before applying Sigmoid, which speeds-up the training for the gates, leading to better results.
A. T RAINING DETAILS
In this section we describe the training procedure for the experiments.
We train all the models with
Adam optimizer (Kingma & Ba, 2014).
2D toy experiments : the models are trained for 500 thousand iterations, with learning rate 1e-4 and cosine annealing.
The mixtures of Gaussians are initialized with means evenly spaced betweenand 4 for EMF-T and between -10 and 10 for EMF-M, while the standard deviations are initialized to 1 For EMF-T and 3 for EMF-M.
MNIST : the models are trained with early stopping until there is no improvement in the validation loss for 100 epochs, with learning rate 1e-4.
The mixtures of Gaussians are initialized with means evenly spaced between -15 and 15 for EMF-T and between -20 and 20 for EMF-M, while the standard deviations are initialized to 3 For EMF-T and 1 for EMF-M.
Generative Hierarchical : we fit the models for 100 thousand iterations for the IRIS dataset and for 400 thousand iterations for the Digits dataset, with learning rate 1e-4 and cosine annealing.
The annealing schedule is over 500 thousands iterations.
Generative SDEs : the models are trained for 50 thousand iterations on BR and OU, while forthousand iterations for LZ and VDP, with learning rate 1e-4 and cosine annealing.
The annealing schedule is over 500 thousand iterations.
Variational inference : we fit all the surrogate posteriors for 100000 iterations with full-batch gradient descent, using a 50-samples Monte Carlo estimate of the ELBO, and learning rate 1e-3 for all the methods but IAF, EMF-T and GEMF-T, for which we use 5e-5.
For GEMF, the gates are initialized to be close to the prior, whit a value of 0.999 after Sigmoid activation.
A. N UMBER OF P ARAMETERS AND SAMPLE TIME
We report the number of trainable parameters and sample time in seconds for each model used in the generative experiments (tables 4, 5, 6, 7, 8, 9).
Note that the samples were performed on a GPU
Nvidia Quadro RTX 6000 .
We do not report parameters and inference time for the Variational
Inference experiments, as the number of parameter changes only with gates, in which there is one additional parameter per latent variable, and the inference time difference between models with and without EMF layer is similar to the sampling time in the generative case
[Table]
[Table] [Table] [Table]
[Table] [Table]
inversion needs a bigger amount of GPU memory
[Table]
A. N UMERICAL ROOT FINDING METHODS
There are cases in which a closed form inversion for the transformation induced by EMF is not available, such as for the Mixture of Gaussian case.
In those cases, we can still compute the inverse function and log-Jacobian determinants by using numerical root finding methods.
In this work, we use either the secant method or the Chandrupatla’s method (Chandrupatla, 1997; Scherer, 2010), as the implementation is available in the probabilistic programming framework Tensorflow Probability.
A.6. S ECANT METHOD
The secant method is a root-finding algorithm that uses a succession of roots of secant lines to better approximate a root of a function.
The secant method can be thought of as a finite-difference approximation of Newton’s method.
The secant method starts with two initial values x 0 and x 1 which should be chosen to lie close to the root, and then uses the following recurrent relation:
[Equation]
A.6. C HANDRUPATLA ’ S METHOD
This root-finding algorithm is guaranteed to converge if a root lies within the given bounds.
At each step, it performs either bisection or inverse quadratic interpolation.
The specific procedure can be found in (Chandrupatla, 1997; Scherer, 2010).
[Figure]
[Figure]
A. A DDITIONAL RESULTS A.7. D ENSITY PLOTS FOR 2D TOY EXPERIMENTS
We include density plots for the 2D toy experiments in figures 3 and 4.
A.7. MNIST SAMPLES
We include samples for models trained on MNIST in figure 5.
A.7. H IERARCHICAL MODELS
Hierarchical models are central to both Bayesian and frequentist statistics and can be used to model data with nested sub-populations. For example, a Gaussian hierarchical model with m populations each with n datapoints can be expressed as [Equation]
where m k denotes the empirical population mean.
For the leaf variables, this results in the following structured layer: f jk,σ ( x jk , η j ) =
ση k + νx jk , where x jk is the latent variable of the leaf node x jk while η k is the latent of its root node.
This extremely simple layer can be used to couple variables that are known to be hierarchically related and can be straightforwardly generalized to deeper and more complex hierarchical couplings.
Again, the downstream and upstream layers in the EMF architecture can model additional couplings and non-Gaussianity while keeping the hierarchical structure.
To test the effectiveness of EMF with Gaussian hierarchical structure, we train the flows on a hierarchical generative modeling problem where each network outputs a whole dataset subdivided into several classes.
Specifically, we use two classification datasets, the IRIS dataset, in which three flower classes are described by four numerical features, and the Digits dataset, composed of 8 × 8 images of handwritten digits from 0 to 9.
We modify the datasets by using vectors ¯ x = [¯ µ, ¯ d 1 , . . .
, ¯ d n ] , where the ¯ d i , i ∈ [1 , . .
. , n ] are datapoints randomly sampled from the same class, and ¯ µ is the feature-wise mean of such datapoints.
We then train the flows with vectors ¯ x ′ =
[¯ µ, ¯ d 1 , . . .
, ¯ d n − 1 ] , as the last datapoint ¯ d n can be inferred from the mean and the other n − 1 datapoints.
We use the same MAF and
MAF-L baselines as for the 2d toy problems, and the gated version of EMF-T and EMF-M (GEMF-T,
GEMF-M).
The results are reported in table 10.
The EMF architectures perform better than their
MAF equivalents but worse than the much more highly parameterized MAF-L architectures.
This is not surprising given the extreme simplicity of the hierarchical layer.
[Table]
A.7. S MOOTHING STRUCTURE FOR TIME SERIES
We can induce smoothness using a layer obtained by discretizing the second-order equation ¨ x ( t ) = − a ˙ x ( t ) + w ( t ) [Table] which models a physical Brownian motion.
We show the results obtained in table 11.
[Table] [Table]
VI EXPERIMENTS
As additional variational inference experiments
[Table]
we show the results on the time series with Gaussian (regression) emissions and a shallow version of the Gaussian binary tree model with depth 4.
The results are shown in tables 12.
We also compute the forward KL divergence for each model.
Such results are reported in table [Table]
The EMF bijective transformation can be also used in combination with models which do not use normalizing flows.
We combine EMF with the mean field and multivariate normal approximations from (Kucukelbir et al., 2017), with results in table 14.
We also report results of the non gated version of EMF.
[Table] [Table]
[Table]
A.7. S URROGATE POSTERIORS
We show and compare the obtained surrogate posteriors using IAF, EMF-T and GEMF-T, together with the ground truth and the observations (figures 6, 7, 8, 9).
A.7. T RAINING LOSSES
We show a comparison of the training losses for the different models, for the 2D data (figure 10,
MMIST (figure 11, the IRIS and digits datsests (figure 13) and for the time series data (figure 12).
A.7. S AMPLES
We provide a comparison between the training data and the samples from the trained models.
Figures 14 15 and 16 show samples for the Brownian motion, Ornstein-Uhlenbeck process and Lorenz system respectively.