Embedded-model flows: Combining the inductive biases of model-free deep learning and explicit probabilistic modeling
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 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 state-of-the-art methods in common structured inference problems.
1. Introduction
[Figure]
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).
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 domainspecific 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.
2. Converting 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]
In practice the inputs 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.
Appendix A.3 describes the single variable transformation, while the next section presents a simple approach to construct structured bijections (i.e., flows) from multivariate distributions expressed as probabilistic programs.
2.1. From differentiable probabilistic programs to structured layers
A probabilistic program samples a sequence of variables in turn, where (as in eq.
(8)) the distribution of each variable is conditioned on the previously-sampled variables.
At the sampling site of the j -th random variable x j , we therefore have all the information necessary to represent its conditional distribution as an invertible transformation f j, ϕ ( ϵ 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 ) (pseudocode in A.4).
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:
[Equation]
In the reverse direction, the variables x 1: n are given, so all dimensions 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 j ; x <j ) is the Jacobian of the local inverse transformation f − 1 ϕ ( x j ; x <j ).
When used in a normalizing flow, we denote the transformation F ϕ
( ϵ 1: n ) as a structured layer .
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]
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 ϕ ( x j ; x <j , λ 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.
(2), eq.
(3), and eq.
(4) 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.
2.2. Automation
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. 8 is recursively converted into a bijective transformation as shown in the 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.
We will release the full TensorFlow Probability code for automatic structured layer construction prior to publication.
2.3. Embedded-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.
3. Related 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 and 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 gauge-equivariant 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 and 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 and 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 and 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.
4. Applications
Model-embedding allows users of all levels of expertise to insert domain knowledge into normalizing flow architectures.
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.5, A.6 and A.7 respectively.
4.1. Variational inference
[Figure]
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 (GEMF-T) with a two-layers inverse autoregressive flow (IAF) and a final structured layer embedding the prior model.
This is visualized in figure 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) and a Lorenz system (LZ), 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 1.
Results on additional experiments can be found in Appendix A.9.1.
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.
[Table]
5. Discussion
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 we achieved state-of-the-art results on a series of structured variational inference problems.
Limitations: 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.
Another limitation is that, 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.
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.
Appendix A. Appendix A.1. Preliminaries
Normalizing flows define a random variable as an invertible differentiable transformation x = f ϕ ( ϵ ) of a base variable ϵ ∼ p 0 ( ϵ ).
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,
[Equation]
where J ϕ ( x ) is the Jacobian matrix of f ϕ .
The base distribution p 0 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 and 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).
A.2. Probabilistic programs and 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 logdensity 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
i 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 x
[Equation]
[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.
A.3. Individual random variables as bijective transformations
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 ) = (cid:82) ya p ϕ ( y )
d y , 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.8), 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.
(7) only involves f − 1 ϕ and its derivatives.
Example: Gaussian Variables
The inverse CDF of a Gaussian variable x ∼ 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 µ,σ ( ϵ ) =
Φ − 1 µ,σ (Φ 0 , 1 ( ϵ ))
= µ + σϵ .
This is the famous reparameterization formula used to train VAEs and other variational methods (Kingma and 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.
A.4. Pseudocode
The pseudocode is shown in figure A.4. A.5. Datasets
Here we provide a detailed description of the datasets used in the experiments.
A.5.1. Stochastic differential equations
We generate data using two stochastic differential equations (SDEs) models, discretized with the Euler-Maruyama method.
The SDEs are: Brownian motion and Lorenz system.
In the following, W t corresponds to a Wiener process, and the series length T = 30 unless specified.
def forward and log det jacobian( model generator, epsilons, gated=
True): x, xs =
None, [] forward log det jacobian =gen = model generator() d = gen.send(None)
# Dist.
of first RV.
for eps in epsilons: # Apply the j’th local transformation.
x, fldj term = ( d.as bijector(gated=gated).
forward and log det jacobian(eps))
xs.append(x) forward log det jacobian += fldj term # Run the model to the next sample site.
d = gen.send(x)
# Needs output ‘x‘.
return xs, forward log det jacobian def inverse and log det jacobian( model generator, xs, gated=
True): epsilons =
[] inverse log det jacobian =gen = model generator() d = gen.send(None)
# Dist.
of first RV.
for x in xs:
# Invert the j’th local transformation.
eps, ildj term = ( d.as bijector(gated=gated).
inverse and log det jacobian(x)) epsilons.append(eps) inverse log det jacobian += ildj term # Run the model to the next sample site.
d = gen.send(x)
# ‘eps‘ is unused.
return epsilons, inverse log det jacobian
Brownian motion : the dynamics evolve as: [Equation] where µ and σ are constants. The Brownian motion sequence is generated as:
[Equation]
where the mean µ = 0 and the standard deviation σ = 0 . 1.
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 .
A.5.2. Smoothing and Bridge data
The data is generated with either a Brownian mmotion or a Lorenz system.
In both cases, 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 and σ e = 1 for the
Lorenz system.
In the classification case, the emissions follow a Bernoulli distribution:
[Equation]
with k a gain parameter, k = 5 in the Brownian Motion and k = 2 in the Lorenz system.
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.
For the Lorenz system, the emissions are observed only for the x variables in both smoothing and bridge.
A.5.3. Hierarchical Models
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.6. Flow Models
In all the experiments, we use IAF with two autoregressive layers 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 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.
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.7. Training details
In this section we describe the training procedure for the experiments.
We train all the models with Adam optimizer (Kingma and Ba, 2014).
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.
A.8. Numerical 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.8.1. Secant 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.8.2. Chandrupatla’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).
A.9. Additional results A.9.1. 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 2.
We also compute the forward KL divergence for each model.
Such results are reported in table 3.
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 4.
We also report results of the non gated version of EMF.
[Table]
[Table]
[Table]
A.9.2. Surrogate posteriors
We show and compare the obtained surrogate posteriors using IAF, EMF-T and GMEF-T, together with the ground truth and the observations (figures 4 and 5).
[Figure] [Figure]