Automatic Variational Inference with Cascading Flows
Abstract
The automation of probabilistic reasoning is one of the primary aims of machine learning.
Recently, the conﬂuence of variational inference and deep learning has led to powerful and ﬂexible automatic inference methods that can be trained by stochastic gradient descent.
In particular, normalizing ﬂows are highly parameterized deep models that can ﬁt arbitrarily complex posterior densities.
However, normalizing ﬂows struggle in highly structured probabilistic programs as they need to relearn the forward-pass of the program.
Automatic structured variational inference (ASVI) remedies this problem by constructing variational programs that embed the forward-pass.
Here, we combine the ﬂexibility of normalizing ﬂows and the prior-embedding property of ASVI in a new family of variational programs, which we named cascading ﬂows.
A cascading ﬂows program interposes a newly designed highway ﬂow architecture in between the conditional distributions of the prior program such as to steer it toward the observed data.
These programs can be constructed automatically from an input probabilistic program and can also be amortized automatically.
We evaluate the performance of the new variational programs in a series of structured inference problems.
We ﬁnd that cascading ﬂows have much higher performance than both normalizing ﬂows and ASVI in a large set of structured inference problems.
1. Introduction
The aim of probabilistic programming is to provide a fully automated software system for statistical inference on arbitrary user-speciﬁed probabilistic models (also referred to as probabilistic programs) (Milch et al., 2007; Sato, 1997;
>. Kersting and De Raedt, 2007; Pfeffer, 2001; Park et al., 2005; Goodman et al., 2012; Wingate et al., 2011;
Patil et al., 2010; Dillon et al., 2017; Bingham et al., 2019; Tran et al., 2017; 2016).
When the probabilistic programs are built from differentiable components, stochastic variational inference (VI) by automatic differentiation offers an effective and computationally efﬁcient solution (Bishop et al., 2002; Kucukelbir et al., 2017; Tran et al., 2017).
However, the performance of stochastic VI depends strongly on the choice of a variational program (also known as variational family or variational guide) (Bishop and Winn, 2003).
Until recently, automatic variational program construction was limited to simple mean ﬁeld and multivariate normal coupling approaches (Kucukelbir et al., 2017).
These programs are highly constrained and exploit only minimally the structure of the original probabilistic program.
For example, the multivariate normal approach only depends on the number of variables and their support.
More recently, normalizing ﬂows revolutionized the ﬁeld of variational inference by offering a highly ﬂexible parametric model for complex multivariate distributions.
Normalizing ﬂows can be easily used to construct automatic variational programs by mapping all the variables to a spherical normal latent space through a learnable diffeomorphism (Rezende and Mohamed, 2015; Papamakarios et al., 2019).
While highly ﬂexible, the normalizing ﬂow approach still uses only minimal information concerning the input probabilistic program.
This leads to poor performance in highly structured probabilistic programs with complex computational ﬂows (Ambrogioni et al., 2021).
For example, consider an inference problem deﬁned by a highly parameterized natural language model (Brown et al., 2020).
A typical problem of this kind is controlled language generation, where the goal is to steer the model to general text conditioned on its past output and future targets (e.g. the presence of some speciﬁc words).
The controlled network can be formulated in probabilistic programming terms as a posterior distribution that is parameterized by the variational program.
In a problem like this, most of the complexity of the posterior is encoded in the structure of the prior probabilistic program (the language model).
The model just needs to be steered in the right direction.
However, normalizing ﬂow completely disregards this structure and has to relearn all the structure embedded in the prior.
This problem was discussed in a recent work that introduced a form of automatic structured variational inference (ASVI) with an automated variational program that incorporates the prior probabilistic program as a special case.
ASVI has been shown to have very good performance in time series analysis and deep learning problems in spite of being very parsimoniously parameterized compared to normalizing ﬂows (Ambrogioni et al., 2021).
However, the convex-update families used in ASVI have two main limitations: I) they cannot model statistical dependencies arising from colliders in the graphical model (the "explaining away" effect) and II) they constrain the conditional posterior distributions to be in the same family as the prior.
In this paper we integrate ASVI and normalizing ﬂows to obtain a new form of automatic VI, which we refer to as cascading ﬂows (CF).
The approach incorporates the forward-pass of the probabilistic program while being capable of modeling collider dependencies and arbitrarily complex conditional distributions.
To this end, we make use of a novel highway ﬂow architecture.
CF respects the following design principles
: • Automation: It can be constructed using a fully automatic algorithm that takes the prior probabilistic program as input.
• Locality: It is constructed by locally transforming each conditional distribution in the conditional prior so that it inherits the graphical structure of the input program (with potentially some extra coupling).
This ensures scalability and modularity.
• Prior information ﬂow: It embeds the forward-pass of the probabilistic program.
2. Related work
Automatic algorithms for VI date back to VIBES (Bishop et al., 2002; Bishop and Winn, 2003).
These older works are based on self-consistency equations and variational message passing (Winn and Bishop, 2005).
The use of gradientbased automatic VI for probabilistic programming was pioneered in (Wingate and Weber, 2013) using a mean ﬁeld approach.
Automatic differentiation VI (ADVI) perfected this approach by exploiting the automatic-differentiation tools developed in the deep learning community (Kucukelbir et al., 2017).
ADVI uses ﬁxed invertible transformations to map each variable into an unbounded domain where the density can be modeled with either univariate of multivariate normal distributions.
Structured stochastic VI (Hoffman and Blei, 2015) exploits the local conjugate update rules of locally conjugate models.
This has the major beneﬁt of preserving the local structure of the probabilistic program but it severely limits the class of possible input programs since local conjugacy is a stringent condition.
Automatic structured VI (ASVI) lifts this constraint as it applies a trainable convex-update rule to non-conjugate links (Ambrogioni et al., 2021).
Several other papers introduced new techniques for constructing structured variational programs.
Copula VI introduces the use of vine copulas to model the statistical coupling between the latent variables in the program (Tran et al., 2015).
Similarly, hierarchical VI uses auxiliary latent variables to achieve the same goal (Ranganath et al., 2016).
Neither of those approaches are truly automatic however, as they do not prescribe a unique way to construct the coupling functions.
Normalizing ﬂows are very expressive and highly parameterized models that use deep learning components to approximate complex multivariate densities (Rezende and Mohamed, 2015; Dinh et al., 2014; 2017).
Normalizing ﬂows can approximate arbitrarily complex densities and can therefore be turned into powerful automatic VI methods (Rezende and Mohamed, 2015; Kingma et al., 2016).
However, conventional normalizing ﬂows do not exploit the input probabilistic program.
Structured conditional continuous normalizing ﬂows are a newly introduced class of ﬂows constrained to have the same conditional independence structure of the true posterior (Weilbach et al., 2020).
However, these ﬂow architectures only inherit the graphical structure of the input program, without incorporating the forward pass (i.e. the form of the conditional distributions and link functions).
Modern deep probabilistic programming frameworks provide automatic implementation of many of variational programs.
For example, PyMC (Patil et al., 2010) provides an implementation of ADVI, Pyro (Bingham et al., 2019) implements the automatic construction of mean ﬁeld, multivariate normal and normalizing ﬂow programs (i.e. variational guides) and TensorFlow probability also offers an automatic implementation of ASVI (Dillon et al., 2017; Ambrogioni et al., 2021).
Our new highway ﬂow architecture is a key component of our new approach and it is inspired by highway networks (Srivastava et al., 2015).
To the best of our knowledge, highway ﬂows are the ﬁrst form of residual-like ﬂows with a tractable Jacobian determinant that can be expressed in closed form.
Existing residual ﬂows express the log determinant of the Jacobian as an inﬁnite series that needs to be estimated using a “Russian roulette” estimator (Chen et al., 2019).
We achieve tractability by applying a highway operator to each single layer of the network.
The use of auxiliary variables has proved to be an effective way to improve the expressiveness of Normalizing Flows by operating in an augmented space (Huang et al., 2020; Cornish et al., 2020; Caterini et al., 2020).
Our use of auxiliary variables is similar to the approach adopted in (Dupont et al., 2019) and (Weilbach et al., 2020).
The bound we use for training the augmented model was derived in (Ranganath et al., 2016).
Structured variational programs have the most clear beneﬁts in timeseries analysis problems (Eddy, 1996; Foti et al., 2014; Johnson and Willsky, 2014; Karl et al., 2016;
Fortunato et al., 2017).
Our approach, together with ASVI, differs from the conventional timeseries approach by exploiting the structure of both temporal and non-temporal variables in the model.
3. Preliminaries
In this section we introduce the notation used in the rest of the paper and discuss in detail the methods that form the basis for our work.
3.1. Differentiable probabilistic programs
We denote the values of (potentially multivariate) random variables using lowercase notation and arrays of random variables using boldface notation.
We consider probabilistic programs constructed as a chain of conditional probabilities over an array of random variables x
= ( x 1 , . .
, x N ) .
We denote the array of values of the parents of the j -th variable as π j ⊆ { x i } i (cid:54) =
j , which is a sub-array of parent variables such that the resulting graphical model is a directed acyclic graph.
Using this notation, we can express the joint density of a probabilistic program as follows:
[Equation]
where θ j ( π j ) is a link function that maps the values of the parents of the j -th variable to the parameters of its density ρ j ( · | · ) .
For example, if the density is Gaussian the link function outputs the value of its mean and scale given the values of the parents.
Note that θ j is a constant when the array of parents is empty.
3.2. Convex-update variational programs
An automatic structured variational inference method provides an algorithm that takes a probabilistic program as input and outputs a variational program q ( x ) .
Convex-update variational programs (Ambrogioni et al., 2021) are deﬁned by the following transformation:
[Equation]
with convex-update operator
[Equation]
where λ j and α j are (potentially vector-valued) learnable parameters and (cid:12) is the element-wise product.
This reduces to the prior probabilistic program for λ j → 0 and to a mean ﬁeld variational program for λ j → 1 .
Convex-update ASVI has the advantage of preserving the forward-pass of the probabilistic program, which is often a major source of statistical coupling in the posterior distribution.
The variation program can be trained to ﬁt the posterior by minimizing the evidence lower bound (ELBO) by stochastic gradient descent (SGD) on the differentiable parameters λ j and α j .
3.3. Variational inference with normalizing ﬂows
Normalizing ﬂows express the density of arbitrary multivariate distributions as the push forward of a known base distribution through a differentiable invertible function (i.e. a diffeomorphism).
The base distribution is usually a spherical normal distribution, hence the "normalizing" in the name.
Consider a d -variate probability density p 0 ( z ) and a diffeomorphism Ψ : R d → R d .
We can derive the density of the new random variable x = Ψ( z ) using the change of variable formula
[Equation]
where Ψ − 1 is the inverse transformation and J ( z ) is the
Jacobi matrix of Ψ .
We can write this more compactly using the push-forward operator T Ψ associated to Ψ that maps the density p 0 ( z ) to the density p X ( x ) given in Eq. 4:
[Equation]
Now consider a ﬂexible family of functions Ψ w ( x ) parameterized by the "weights" w .
We can use Eq.
13 to approximate arbitrarily complicated target densities by training w by SGD.
Given a probabilistic program p ( x, y ) with y observed, we can approximate the posterior p ( x | y ) by minimizing the evidence lower bound:
[Equation]
where the expectation is taken with respect to the transformed distribution T Ψ w [ p 0 ( · )] ( x ) .
Since Ψ is differentiable, we can obtain the reparameterization gradient for the
ELBO as follows:
[Equation]
where the expectation is now taken with respect to the ﬁxed distribution p 0 ( z ) .
In the following we refer to this approach as global ﬂow (GF) since the transformation is applied synchronously to all variables in the model.
4. Variational inference with cascading ﬂows
Now that we have covered the relevant background, we will introduce our new approach that combines ASVI with normalizing ﬂows.
4.1. Cascading ﬂows variational programs
The convex-update parameterization poses strong constraints on the form of the conditional posteriors.
On the other hand, normalizing ﬂow programs do not embed the potentially very complex structure of the input program.
In this paper we combine the two approaches by replacing the convex-update operator in Eq.
2 with the push-forward operator used in normalizing ﬂows.
More precisely, consider a family of diffeomorphisms Ψ w =
(Ψ w 1 1 , . . .
, Ψ w N N )
depending on the array of parameters w = ( w 1 , . . .
, w N ) .
Using these transformations, we can deﬁne the cascading ﬂows variational program associated with a program deﬁned by Eq. 1:
[Equation]
where we denoted the push-forward operator induced by the diffeomorphism Ψ w j j (as deﬁned by Eq. 4) as T w j to simplify the notation.
A cascading ﬂows variational program can be trained by SGD using gradient estimators of the ELBO with the reparameterization given in Eq. 6.
The name cascading ﬂows comes from the fact that, since the transformation is performed locally for each conditional distribution, its effects cascade down the variational model through the (transformed) conditional dependencies of the prior program.
However, in order to preserve the forwardpass of the probabilistic program it is crucial to use a parameterized diffeomorphism that can be initialized around the identity function and whose deviation from the identity can be easily controlled.
Speciﬁcally, we opt for transformations of the following form:
[Equation]
where γ ∈ (0 , 1) .
The resulting architecture of a cascading ﬂows variational program is visualized in Fig 1.
From the diagram it is clear that a cascading ﬂow is constructed by inserting transformation layers within the architecture of the probabilistic program.
Transformations of the form given in Eq.
8 with tractable Jacobians are not currently present in the normalizing ﬂow literature.
Therefore in the next subsection we will introduce a new ﬂow architecture of this form, which we named highway ﬂow .
4.2. Highway ﬂow architecture
Here we introduce highway ﬂow blocks gated by a learnable scalar gate parameter.
As also visualized in Fig.
2, a highway ﬂow block is comprised by three separate highway layers:
[Figure]
j
[Figure]
Layer j
+
Layer j Layer j
+
Layer j
+
CF Layer j CF Layer j
+
[Figure] [Figure]
Upper triangular highway layer:
[Equation]
Lower triangular layer:
[Equation]
Highway activation functions:
[Equation]
where U is an upper-triangular matrix with positive-valued diagonal, L is a lower-triangular matrix with ones in the diagonal and g ( x ) is a differentiable non-decreasing activation function.
A highway ﬂow layer is a composition of these three types of layers:
[Equation]
A highway network is a sequence of M highway layers with a common gate parameter and different weights and biases.
Note that a highway ﬂow can be expressed in the form given in Eq. 8 by deﬁning γ as follows:
[Equation]
which is clearly equal to one (zero) when λ is equal to one (zero).
4.3. Auxiliary variables and inﬁnite mixtures of ﬂows
The expressiveness of a normalizing ﬂow is limited by the invertible nature of the transformation.
This is particularly problematic for low-dimensional variables since the expressivity of many ﬂow architectures depends on the dimensionality of the input (Papamakarios et al., 2019).
This limitation can be a major shortcoming in cascading ﬂows as they are particularly useful in highly structured problems characterized by a large number of coupled low-dimensional variables.
Fortunately, we can remedy this limitation by introducing a set of auxiliary variables.
For each variable x j , we create a D -dimensional variable (cid:15)
j following a base distribution p j ( (cid:15) j ) .
We can now use an augmented diffeomorphism ˆΨ w j j ( x j , (cid:15) j ) and deﬁne the joint posterior over both x j and (cid:15) j :
[Equation]
where the push-forward operator now transforms the (independent) joint density of x j and (cid:15) j .
The conditional variational posterior is then obtained by marginalizing out (cid:15)
j :
[Equation]
This inﬁnite mixture of ﬂows is much more capable of modeling complex and potentially multi-modal distributions.
The use of latent mixture models for VI was originally introduced in (Ranganath et al., 2016).
Given a set of observations y , the ELBO of any mixture variational probabilistic program can be lower bounded using Jensen’s inequality:
[Equation]
[Equation]
where q ( (cid:15) )
= (cid:82) q ( x , (cid:15) )
d x is the marginal variational posterior and r ( (cid:15) ) is an arbitrarily chosen distribution over the auxiliary variables.
This result justiﬁes the use of the augmented ELBO for training the variational posterior.
In practice, the distribution r ( (cid:15) ) can be chosen to be a mean ﬁeld variational density.
The gap between the mixture ELBO (Eq. 19) and the augmented
ELBO can be reduced by training these mean ﬁeld parameters by minimizing the augmented ELBO such as to match q ( (cid:15) ) (Ranganath et al., 2016).
The highway ﬂow architecture can be easily adapted to the augmented variable space.
Since there is no need to gate the ﬂow of the auxiliary variables, the scalar λ in Eq.
9, Eq.and Eq.
13 should be replaced by the vector l whose entries corresponding to the original variables are equal to λ while the entries corresponding to the auxiliary variables are equal to zero.
4.4. Modeling collider dependencies with auxiliary coupling
[Equation]
So far, we outlined a model that has the same conditional independence structure of the prior and is therefore incapable of modeling dependencies arising from colliders.
Fortunately, the local auxiliary variables that we introduced in order to increase the ﬂexibility of the local ﬂows can assume a second role as "input ports" that can induce non-local couplings.
In fact, by coupling their auxiliary variables we can induce statistical dependencies between any subset of variables in the model.
Importantly, this coupling can take a very simple form since additional complexity can be modeled by the highway ﬂows.
We will now introduce a local auxiliary model with a non-trivial conditional independence structure inspired by the forward-backward algorithm.
The graphical structure of the posterior probabilistic program is a subgraph of the moralization of the prior graph (Bishop, 2006).
Therefore, in order to be able to capture all the dependencies in the posterior, it sufﬁces to couple the auxiliary variables by reversing the arrows of the original prior model.
In fact, all the parents of every node get coupled by the reversed arrows.
As an example, consider the following probabilistic model (Fig.
3A):
[Equation]
We couple the auxiliary variables by reversing the arrows, as [Table] formalized in the graphical model p 1 ( (cid:15) 1 ) p 2 ( (cid:15) 2 , | (cid:15) 1 ) p 3 ( (cid:15) 3 , | (cid:15) 1 ) p 4 ( (cid:15) 4 | (cid:15) 3 ) .
The graphical structure of the resulting augmented model is shown in Fig.
3B.
The auxiliary variables can be coupled in a very simple manner since the diffeomorphisms can add arbitrary complexity to the ﬁnal joint density.
Here we use a simple linear normal model:
[Equation]
where a j ≥ 0 and (cid:80) n a n = 1 .
In this expression, ξ k is a standard Gaussian vector and υ k
=
( υ 1 , ..., υ K ) is an array containing all the parents auxiliary variables (the auxiliaries of the children of x k in the original graph).
4.5. Inference amortization
We will now discuss a strategy to implement automatic inference amortization in cascading ﬂows models.
The problem of inference amortization is closely connected to the problem of modeling collider dependencies as the latter arise from the backward messages originating from observed nodes (Bishop, 2006).
It is therefore not a coincidence that we can implement both collider dependencies and amortization by shifting the statistics of the local auxiliary variables.
This results in a very parsimonious amortization scheme where data streams are fed locally to the various nodes and propagate backward through the linear network of auxiliary Gaussian variables.
We denote the set of auxiliary variables associated to the children of x k in the input probabilistic program as υ k .
Furthermore, we denote the observed value of x k as y k and the set of observed values of the j -th child of x k as y k
=
( y k 1 , ..., y kK ) .
The [Figure] algorithm. amortized auxiliary model can then be obtained as a modiﬁcation of Eq.
20: p 0 ( (cid:15) k | υ k ) =
N (cid:0) m k , σ 2 I (cid:1) with (cid:15) k
| υ k = B ( k )
[ y k ] + (cid:80) Kj =1 a j (cid:12) υ j + a 0 (cid:12) ξ k , where
B ( k ) are learnable afﬁne transformations (i.e. linear layers in deep learning lingo).
Cascading ﬂow variational programs can be constructed automatically given an input probabilistic program. We provide an open-source implementation of this algorithm in TensorFlow Probability (Dillon et al., 2017). The code for the spacial case without amortization and backward auxiliary coupling is shown in Figure 5.
[Figure]
6. Experiments
Our experiments are divided into three sections.
In the ﬁrst section we focus on highly structured timeseries problems exhibiting linear and non-linear dynamics.
Bayesian problems of this nature have relevance for several scientiﬁc ﬁelds such as physics, chemistry, biology, neuroscience and population dynamics.
In the second section we compare the capacity of several variational programs to model dependencies arising from colliders in the graphical structure.
In fact, the inability of modeling collider dependencies was the main limitation of the original convex-update ASVI family (Ambrogioni et al., 2021).
Finally, in the third section we test the performance of the automatic amortization scheme.
All the code used in the experiments is included in the supplementary material and documented in Supplementary
A. Architectures :
In all experiments, the CF architectures were comprised of three highway ﬂow blocks with softplus activation functions in each block except for the last which had linear activations.
CF programs use an independent network of this form for each latent variable in the program.
Each variable was supplemented with 10 auxiliary variables, the width of each network was therefore equal to the dimensionality of the variable plus 10 .
Weights and biases were initialized from centered normal distributions with scale 0 .
The λ variable was deﬁned independently for each network as the logistic sigmoid of a learnable parameter l , which was initialized as 4 in order to keep the variational program close to the input program.
For each variable, the posterior auxiliary distributions r ( (cid:15) ) (see Eq. 19) were spherical normals parameterized by mean and standard deviation (i.e. Gaussian mean ﬁeld).
[Figure]
6.1. Timeseries analysis
We consider discretized SDEs, deﬁned by conditional densities of the form:
[Equation]
where µ ( x t , t ) is the drift function and σ ( x t , t ) is the volatility function.
Speciﬁcally, we used Brownian motions (BR: µ ( x ) = x , σ 2 ( x ) = 1 ), Lorentz dynamical systems (LZ: µ ( x 1 , x 2 , x 3 ) =
(10( x 2 − x 1 ) , x 1 (28 − x 3 ) − x 2 , x 1 x 2 − 8 / 3 x 3 ) , σ 2 = 2 ), Lotka-Volterra population dynamics models (PD: µ ( x 1 , x 2 ) =
2 x 1 − 0 .
02 x 1 x 2 , 0 .
1 x 1 x 2 − 0 .
1 x 2 ) , σ 2 = 3 ) and recurrent neural dynamics with randomized weights (RNN, see Supplementary B).
The initial distributions were spherical Gaussians (see Supplementary
B for the parameters).
For each dynamical model we generated noisy observations using Gaussian emission models (denoted with an "r" in the tables):
[Equation]
or Bernoulli logistic models ("c" in the tables):
[Equation]
where [ x t ] 1 denotes the ﬁrst entry of the vector x t , σ lk is the standard deviation of the Gaussian emission model, [Table] g ( · ) is the logistic sigmoid function and k is a gain factor.
All the numeric values of the parameters in different experiments are given in Supplementary B.
For each task we evaluate the performance of the cascading ﬂows programs (CF) against a suite of baseline variational posteriors, including convex-update ASVI (Ambrogioni et al., 2021), mean ﬁeld (MF) and multivariate normal (MN)
(Kucukelbir et al., 2017).
We also compare our new cascading ﬂows family with the more conventional global normalizing ﬂow (GF) approach where a single diffeomorphism transforms all the latent variables in the probabilistic program into a spherical normal distribution (Rezende and Mohamed, 2015).
For the sake of comparison, we adopt an architecture identical to our highway ﬂow model (includingauxiliary variables) except for the absence of the highway gates.
Note however that the global architecture has much greater width as it is applied to all variables in the program.
Finally, we compare the performance of CF with a modiﬁed version of CF that does not use highway gates.
Ground-truth multivariate timeseries ˜ x
= (˜ x 1 , . .
˜ x T ) were sampled from the generative model together with simulated ﬁrst-half observations y 1:
T/
=
( y 1 , . . .
, y T/ 2 ) and second-half observations y T/ 2: T = ( y T/ 2+1 , . . .
, y T ) .
All the variational models were trained conditioned only on the ﬁrst half observations.
Performance was assessed using two metrics.
The ﬁrst metric is the average marginal log-probability of the ground-truth given the variational posterior: 1 TJ (cid:80)
Tt (cid:80)
Jj log q ([˜ x t ] j | y 1: T/ 2 ) , where J is the dimensionality of the ground-truth vector ˜ x t .
The marginal posterior density q ([˜ x t ] j | y 1: T/ 2 ) was estimated fromsamples drawn from the variational program using Gaussian kernel density estimation with bandwidth 0 .
9ˆ σN − 1 / 5 , where ˆ σ is the empirical standard deviation of the samples.
Our second metric is log p ( y T/ 2: T | y 1: T/ 2 ) :
the predictive log-probability of the ground-truth observations in the second half of the timeseries given the observations in the ﬁrst half, estimated using kernel density estimation from 5000 predicted observation samples drawn from the exact likelihood conditioned on the samples from the variational program.
Each experiment was repeated 10 times.
In each repetition, all the variational programs were re-trained for 8000 iterations (enough to ensure convergence in all methods) given new sampled ground-truth timeseries and observations.
The reported results are the mean and SEM of the two metrics over these 10 repetitions.
Results:
The results are given in Table 1.
The cascading ﬂows family (CF) outperforms all other approaches by a large margin in all non-linear tasks.
Convex-update ASVI consistently has the second highest performance, easily outperforming the highly parameterized global normalizing ﬂow (GF).
This result is consistent with (Ambrogioni et al., 2021), where
ASVI was shown to outperform inverse autoregressive ﬂow architectures (Kingma et al., 2016).
Finally, the CF model without highway gates (CF non-res) has very low performance, proving the crucial importance of our new highway ﬂow architecture.
A qualitative comparison (LZ-r and LZ-c) between CF and ASVI is shown in Fig.
6.2. Binary tree experiments
In our second experiment we test the capability of the variational programs to model statistical dependencies arising from colliders.
We tested this in a Gaussian binary tree model where the mean of a scalar variable x dj in the d -th layer is a function of two variables in the d − 1 -th layer: x dj ∼ N (cid:0) link ( π d − 1 j 1 , π d − 1 j 2 ) , σ 2 (cid:1) , where π d − 1 j 1 and π d − 1 j 2 are the two parents of x dj .
All the 0 -th layer variables have zero mean.
The inference problem consists of observing the only variable in the last layer and computing the posterior of all the other variables.
We considered a setting with linear coupling link ( x, y ) = x − y and tanh coupling link ( x, y )
=
tanh ( x ) − tanh ( y ) with 2 and 4 layers.
Performance was estimated using the posterior log-probability of the ground-truth sample estimated using a multivariate
Gaussian approximation that matches the empirical mean and covariance matrix of 5000 samples drawn from the trained variational program.
We compared performance with ASVI, mean-ﬁeld (MF), global ﬂow (GF) and multivariate normal (MN).
Note that neither ASVI nor MF can model collider dependencies.
Full details are given in Supplementary C.
The reported results are the mean and SEM of the two metrics over 15 repetitions.
Results:
The results are given in Table 2.
The CF model always outperforms
ASVI and MF, proving that it can capture collider dependencies.
CF has the highest performances in shallow models but lower performance than GF and MN in deeper models.
6.3. Inference amortization
Finally, we evaluate the performance of the amortization scheme in non-linear timeseries inference problems.
We [Table] used the population dynamics (PD) and the Lorentz (LZ) systems.
The CF program was amortized using the approach described in Sec. 4.5.
The MF and normalizing ﬂow baselines were amortized using inference networks (see Supplementary D).
All the details were identical to the timeseries experiments except for the fact that all time points were observed.
Performance was again quantiﬁed as the average marginal log-probability of the ground-truth kernel density estimated from the samples from the amortized variational programs.
Each amortized program was trained once and tested on 50 new sampled timeseries/observations pairs.
The reported results are the mean and SEM of these 50 repetitions.
Results:
The results are reported in Table 3.
Like in the non-amortized case, CF greatly outperforms both MF and GF in both timeseries inference problems.
7. Discussion
Our CF approach solves the two main limitations of ASVI as it can model collider dependencies and can be amortized automatically.
We show that cascading ﬂows programs have extremely high performance in structured inference problems.
However, their ability to capture deep collider dependencies is lower than normalizing ﬂows.
In our opinion, designing principled ways to model collider dependencies is a promising research direction.
[-Diffusion: σ 2 ( x ) = 
-]
[-π d − 1 j 1 − π d − 1 j 
-]