Year: 2021

Abstract:
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, hierarchical coupling 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.

heading: 1 INTRODUCTION
text: 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 data-
generating 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.

heading: 2 PRELIMINARIES
text: Normalizing flows define a random variable as an invertible differentiable transformation x = fϕ(ϵ) of a base variable ϵ ∼ p0(ϵ). In these expressions, ϕ are parameters that control the form of the transformation and consequently the distribution of x. The transformed (log-)density can be given explicitly using the change of variables formula from probability theory,
log p(x) = log p0(f −1 ϕ (x))− log |det Jϕ(x)|, (1)
where Jϕ(x) is the Jacobian matrix of fϕ. The base distribution p0 is often taken to be standard normal (ϵ ∼ N(0, I)). In typical applications, the functional form of fϕ is specified by a deep architecture engineered to guarantee stable and tractable invertibility and tractable (log-)Jacobian computations (Dinh et al., 2014; Rezende & Mohamed, 2015). 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).

heading: 2.1 PROBABILISTIC PROGRAMS AND GRAPHICAL MODELS
text: 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 xi is a graph node connected by incoming edges to some subset of the previous variables x<i. This leads to the following general expression for the joint density of a set of vector-valued variables x1:n = {xj}nj=1:
p(x0:n;ϕ) = p0(x0;θ0(ϕ)) n∏ j=1 pj(xj |θj(x<j ;ϕ)) (2)
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.

heading: 3 CONVERTING PROBABILISTIC PROGRAMS INTO NORMALIZING FLOWS
text: 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:
x = Fϕ(ϵ) ∼ p(x;ϕ). (ϵ ∼ N (0, I)) (3)
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 use 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.

heading: 3.1 INDIVIDUAL RANDOM VARIABLES AS BIJECTIVE TRANSFORMATIONS
text: We first consider the case of a univariate random variable x ∼ pϕ(x), which we can express as a continuous bijective transformation fϕ of a standard normal variable ϵ via a slight extension of inverse transform sampling. The probability integral transform theorem states that the cumulative distribution function (CDF) of pϕ, Cϕ(x) = ∫ x a pϕ(y)dy, 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
x = fϕ(ϵ) = C −1 ϕ (Φ0,1(ϵ)) (ϵ ∼ N (0, 1)) (4)
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 ∼ N ( µ, σ2 ) 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µ,σ(ϵ) = Φ−1µ,σ(Φ0,1(ϵ)) = µ + σϵ. This is the famous reparameterization formula used to train VAEs and other variational methods (Kingma & Welling, 2013). The inverse function is f−1µ,σ(x) = (x− µ)/σ 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.

heading: 3.2 FROM DIFFERENTIABLE PROBABILISTIC PROGRAMS TO STRUCTURED LAYERS
text: A probabilistic program samples a sequence of variables in turn, where (as in eq. (2)) the distribution of each variable is conditioned on the previously-sampled variables. At the sampling site of the j-th random variable xj , we therefore have all the information necessary to represent its conditional distribution as an invertible transformation fj,ϕ(ϵj ;x<j) of a standard normal input ϵj . Feeding the result back into the model defines a sequence of local transformations, which together form a joint flow layer Fϕ(ϵ1:n) (algorithm 1). This layer transforms a spherical Gaussian variable ϵ1:n into a variable that follows the same distribution of the input program.
As with autoregressive flows, the forward transformation of the j-th input requires us to have computed the preceding outputs x<j = F j ϕ(ϵ<j). This may be expressed recursively as follows:
[Fϕ(ϵ1:n)]j = fj,ϕ(ϵj ;F j ϕ(ϵ<j)), (5)
This recursive expression terminates since the parents have an acyclic-directed-graph structure. For variables without parents, F jϕ reduces to a trivial map between empty sets. In the reverse direction, the variables x1:n are given, so all dimensions may be evaluated in parallel:[
F−1ϕ (x1:n) ] j = f−1j,ϕ(xj ;x<j) (6)
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:
log ∣∣∣det J−1ϕ (x1:n)∣∣∣ = ∑
j
log ∣∣∣det J−1ϕ,j (xj ;x<j)∣∣∣ (7)
where J−1ϕ,j (xj ;x<j) is the Jacobian of the local inverse transformation f −1 ϕ (xj ;x<j). When used in a normalizing flow, we denote the transformation Fϕ(ϵ1:n) as a structured layer.
Algorithm 1 Gated forward transformation: ϕ: Model parameters; {ϵk}nk=1: Input variables sorted respecting parenting structure; {λk}nk=1: Gate parameters; n Number of variables.
procedure GATEDFORWARDSAMPLER({ϵk}nk=1) State log |det J | ← 0 for j in Range(n) do {pi} mj i=1 ← ϵj .parents idx ▷ {pi} mj i=1: indices of parents of j-th variable.
▷ mj : number of parents of j-th variable. if {pi} mj i=1 is not empty then
x̃j ← fj,ϕ(ϵj ; {xpi} mj i=1) xj ← λjx̃j + (1− λj)ϵj log |det J | ← log |det J |+ log (det |Jϕ,j({xk}nk=1) + (1− λj)I|)
else x̃j ← fj,ϕ(ϵj ;ϕ0j ) ▷ ϕ0j : parameters of the j-th root distribution. xj ← λjx̃j + (1− λj)ϵj log |det J | ← log |det J |+ log ( det |Jϕj0,j(xj) + (1− λj)I|
) return {xk}nk=1, log |det J |
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 gj,ϕ to be a convex combination of the original transformation fj,ϕ and the identity mapping:
gj,ϕ(ϵj ;x<j , λj) = λjfj,ϕ(ϵj ;x<j) + (1− λj)ϵj (8)
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 λ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ϕ (xj ;x<j , λj) can in general be computed by numeric root search, but the Gaussian case admits a closed form in both directions:
gµ,σ(ϵ, λ) = (1 + λ(σ − 1))ϵ+ λµ; g−1µ,σ(x, λ) = x− λµ
1 + λ(σ − 1) (9)
Analogues of eq. (5), eq. (6), and eq. (7) then define the gated joint transformation Gλ,ϕ, 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.

heading: 3.3 AUTOMATION
text: 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. 2 is recursively converted into a bijective transformation as shown in the code in Fig. 1. 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. We will release the full TensorFlow Probability code for automatic structured layer construction prior to publication.

heading: 3.4 EMBEDDED-MODEL FLOWS
text: 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. 1.

heading: 4 RELATED WORK
text: 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.

heading: 5 APPLICATIONS
text: 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 multimodal models and moving on to more sophisticated structured models for time series and hierarchical models. 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. In the following, results denoted with a * are reported on a fraction of the total amount of runs, as some runs are currently running. Number of parameters and sampling time for all the models are reported in appendix A.5. All the code use in the experiments is available at https: //anonymous.4open.science/r/EmbeddedModelFlows-9172/README.md.

heading: 5.1 MULTIMODALITY
text: 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:
pρ,µ,σ(x) = ∏ j ∑ k ρj,kN(xj ;µjk, σ2jk) . (10)
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
f−1ρ,µ,σ(x) = Φ −1 0,1  n∑ j=1 ρjΦµ,σ(x)  (11) In itself, the probabilistic program in Eq. 11 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 (on par with NSF for the eight gaussians). However, this comes at the cost of additional sampling time, caused by the numeric inversion of the mixture of Gaussians transformation.

heading: 5.2 HIERARCHICAL MODELS
text: 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
p(m,x) = m∏ k=1 N (mk; 0, σ2) n−1∏ j=1 N ( xjk;mk, ν 2 ) (12)
where mk denotes the empirical population mean. For the leaf variables, this results in the following structured layer: fjk,σ(ϵjk, ηj) = σηk + νϵjk , where ϵjk is the latent variable of the leaf node xjk 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 2. 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.

heading: 5.3 TIMESERIES MODELING
text: 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
xt+1 ∼ N ( xt+1;xt, σ 2 ) , (13)
which gives us the bijective function ft+1,σ(xt+1;xt) = xt + σxt+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 sections, 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. For the models using NSF we scale the time series to have standard deviation 1, which we found empirically to improve stability. We do so using a scaling bijective transformation and computing its Jacobian, making the losses with and without scaling comparable. The results are reported in Table 3. GEMF 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.3.
Table 3: Results in negative log probability on one million test samples for the time-series toy problems, and on a test set of 10000 datapoints. We report mean and standard error of the mean over five different runs.
GEMF-T(c) NSF-GEMF-T(c) MAF MAF-L NSF B-MAF
BR −26.414± 0.0115 – −26.061± 0.0103 −26.121± 0.0099 – −25.905± 0.0100 OU 24.086± 0.0015 – 24.197± 0.0028 24.162± 0.0029 – 24.169± 0.0034 LZ −192.420± 0.4389 −176.710± 0.1445∗ −136.626± 1.0225 −128.837± 2.7377 −115.267± 7.8488∗ −160.440± 3.4641 VDP −551.202± 0.3661 – −484.416± 0.9283 −512.586± 0.5547 – −534.165± 0.3736

heading: 5.4 VARIATIONAL INFERENCE
text: IAF
Variational MEF architecture
Prior model
Figure 2: Diagram of a gated prior-embedding architecture with a sample of transformed variables for a Lorentz dynamics smoothing problem.
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 4. Results on additional experiments can be found in Appendix A.7.4. GEMF-T outperforms the baselines on all the problems, showing the benefits of embedding the prior program for highly structured problems, and the capability of such architecture to model collider dependencies.

heading: 6 DISCUSSION
text: 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. 5, 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, expect in the case of variational inference, this not lead to training slow-downs as the inverse transform is always trivially parallelizeble. 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.

heading: A APPENDIX

heading: A.1 DATASETS
text: Here we provide a detailed description of the datasets used in the experiments.
A.1.1 2D TOY PROBLEMS
The implementation of the 2D toy datasets is taken from the publicly available implementation of FFJORD (Grathwohl et al., 2019).

heading: A.1.2 MNIST
text: 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).

heading: A.1.3 STOCHASTIC DIFFERENTIAL EQUATIONS
text: 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, Wt 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: ẋt = µxtdt+ σxtdWt
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:
x0 ∼ N (µ, σ) xt ∼ N (xt−1, σ),∀t ∈ [1, . . . , T − 1]
where the mean µ = 0 and the standard deviation σ = 0.1.
Ornstein-Uhlenbeck process: This process has the following dynamics: ẋt = −θxtdt+ σdWt
with θ > 0 and σ > 0. The sequences are generated as:
x0 ∼ N (µ, σ0) xt ∼ N (θxt−1, σ),∀t ∈ [1, . . . , T − 1]
with µ0 = 0, σ0 = 5, σ = 0.5 and θ = 0.8.
Lorenz system: the Lorenz dynamics evolve as follows:
ẋt = ϕ(yt − xt) ẏt = xt(ρ− zt)− yt żt = xtyt − βzt
where ϕ = 10, ρ = 28 and β = 83 . In the discrete case we have:
x0, y0, z0 ∼ N (0, 1) ∀t ∈ [1, . . . , T − 1] xt ∼ N (xt−1 + s(ϕ(yt−1 − xt−1)), √ s ∗ σ)
yt ∼ N (yt−1 + s(xt−1(ρ− zt−1)− yt−1), √ s ∗ σ) zt ∼ N (zt−1 + s(xt−1yt−1 − βzt−1), √ s ∗ σ)
with standard deviation σ = 0.1 and step size s = 0.02.
Van der Pol oscillator: The system evolves as:
ẋ = y +Wt
ẏ = µ(1− x2)y − x+Wt In the discrete case, with T = 120 and s = 0.05, we have:
x0, y0 ∼ N (0, 1) ∀t ∈ [1, . . . , T − 1] xt ∼ N (xt−1 + yt−1 ∗ s, √ s ∗ σ)
yt ∼ N (yt−1 + sµ(1− x2t−1)yt−1, √ s ∗ σ)
Here we use µ = 1 and σ = 0.1

heading: A.1.4 SMOOTHING AND BRIDGE DATA
text: 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:
et ∼ N (xt, σe)
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:
et ∼ Bernoulli(kxt)
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 last 40 for 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.

heading: A.1.5 HIERARCHICAL MODELS
text: 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:
µ ∼ N (0, 100) log τ ∼ logN (5, 1)
θi ∼ N (µ, τ2) yi ∼ N (θi, σ2i )
where µ represents the prior average treatment effect, τ controls how much variance there is between schools, i = 1, . . . , 8 is the school index, and yi 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, xdj , are sampled from a Gaussian distribution where the mean is function of two parent variables πd−1j,1 , π d−1 j,2 at the previous layer:
xdj ∼ N (link(πd−1j,1 , π d−1 j,2 ), σ 2)
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.

heading: A.2 STRUCTURES INDUCED WITH EMF
text: 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:
f(x) = µ+ σx
f−1(x) = (x− µ)
σ
while in the gated case, with λ ∈ (0, 1) the transformation becomes:
g(x, λ) = λµ+ x(λ(σ − 1) + 1)
g−1(x, λ) = x− λµ
1 + λ(σ − 1)) 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 xt∀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 µ = 2x ′ 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 ′ 0 and σ = 1.

heading: A.3 FLOW MODELS
text: 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.

heading: A.4 TRAINING DETAILS
text: 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 between -4 and 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 0 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 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 GMEF, the gates are initialized to be close to the prior, whit a value of 0.999 after Sigmoid activation.

heading: A.5 NUMBER OF PARAMETERS AND SAMPLE TIME
text: We report the number of trainable parameters and sample time in seconds for each model used in the generative experiments (tables 5, 6, 7, 8, 9, 10). 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.

heading: A.6 NUMERICAL ROOT FINDING METHODS
text: 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.

heading: A.6.1 SECANT METHOD
text: 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 x0 and x1 which should be chosen to lie close to the root, and then uses the following recurrent relation:
xn = xn−1 − f(xn−1) xn−1 − xn−2
f(xn−1)− f(xn−2)
= xn−2f(xn−1)− xn−1f(xn−2)
f(xn−1)− f(xn−2)

heading: A.6.2 CHANDRUPATLA’S METHOD
text: 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).

heading: A.7 ADDITIONAL RESULTS

heading: A.7.1 DENSITY PLOTS FOR 2D TOY EXPERIMENTS
text: We include density plots for the 2D toy experiments in figures 3 and 4.

heading: A.7.2 MNIST SAMPLES
text: We include density plots for the 2D toy experiments in figure 5.

heading: A.7.3 SMOOTHING STRUCTURE FOR TIME SERIES
text: We can induce smoothness using a layer obtained by discretizing the second-order equation ẍ(t) = −aẋ(t) +w(t), which models a physical Brownian motion. We show the results obtained in table 11.
A.7.4 VI EXPERIMENTS
As additional variational inference experiments, 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 13.
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.

heading: A.7.5 SURROGATE POSTERIORS
text: We show and compare the obtained surrogate posteriors using IAF, EMF-T and GMEF-T, together with the ground truth and the observations (figures 6, 7, 8, 9).

heading: A.7.6 TRAINING LOSSES
text: 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).

heading: A.7.7 SAMPLES
text: 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.

Review summaries:
This paper proposes a new type of normalizing flow for incorporating inductive biases into the model architecture. To this end, the authors introduce a so-called “structure layer”, aiming to transform a spherical Gaussian variable into a pre-defined probabilistic program. A modified version, called “gated structured layer” is proposed in order to skip problematic parts of the model. The model is evaluated on a variety of different datasets and outperforms the masked autoregressive flow baseline.

This paper first proposes to take a probabilistic program and compile it into a sequence of invertible transformations. Applying this invertible transformation to a standard normal distribution produces a sample from the probabilistic program. This "compiled" invertible transformation is then used as a layer with learnable parameters, called a structured layer, in a larger flow model which can also have unstructured layers. The paper argues that this enables the ability to incorporate domain knowledge via these structured layers while enabling deviation from this domain knowledge via the unstructured flow layers. The paper performs experiments on multimodal, hierarchical, time series and variational inference data. 

This paper tries to bridge the concept of the normalizing flow model with the structured layers that encode the domain knowledge (also called inductive bias here). The general idea is to first notice that many existing probablistic programs (such as univariable RV) can be converted to flow under some special $f_{\phi}$.  Furthermore, the authors proposed a gated layer to allow the model to switch between user-specified model and the learned MAF. Extensive experiments in toy multimodality distributions, hierarchical gaussian, timeseries models and variational inferences are conducted to compare the proposed method with other baselines.  

The authors describe how to translate a probabilistic program into a normalising flow layer which maps samples from a unit Gaussian into samples from the probabilistic program. They propose a "gating" mechanism for this layer so that it can learn to interpolate between this transformation and the identity transform. They then show how this layer can be combined with generic normalizing flow layers, and in this way add an inductive bias to a normalising flow without sacrificing its expressivity. In the experiments, they provide various examples using probabilistic programs in this way to construct normalising flows with helpful inductive biases.