Original: V ARIATIONAL A UTOENCODERS WITH J OINTLY O PTIMIZED L ATENT D EPENDENCY S TRUCTURE
Revised: V ARIATIONAL A UTOENCODERS WITH J OINTLY O PTIMIZED L ATENT D EPENDENCY S TRUCTURE

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: We propose a method for learning the dependency structure between latent variables in deep latent variable models.
Revised: We propose a method for learning the dependency structure between latent variables in deep latent variable models.

Original: Our general modeling and inference frame
Revised: Our general modeling and inference frame

Original: - work combines the complementary strengths of deep generative models and probabilistic graphical models.
Revised: - work combines the complementary strengths of deep generative models and probabilistic graphical models.

Original: In particular, we express the latent variable space of a variational autoencoder (VAE) in terms of a Bayesian network with a learned, ﬂexible dependency structure.
Revised: In particular, we express the latent variable space of a variational autoencoder (VAE) in terms of a Bayesian network with a learned, ﬂexible dependency structure.

Original: The network parameters, variational parameters as well as the latent topology are optimized simultaneously with a single variational objective.
Revised: The network parameters, variational parameters as well as the latent topology are optimized simultaneously with a single variational objective.

Original: Inference is formulated via a sampling procedure that produces expectations over latent variable structures and incorporates top-down and bottom-up reasoning over latent variable values.
Revised: Inference is formulated via a sampling procedure that produces expectations over latent variable structures and incorporates top-down and bottom-up reasoning over latent variable values.

Original: We validate our framework in extensive experiments on MNIST, Omniglot, and CIFAR-10.
Revised: We validate our framework in extensive experiments on MNIST, Omniglot, and CIFAR-10.

Original: Comparisons to state-of-the-art structured variational autoencoder baselines show improvements in terms of the expressiveness of the learned model.
Revised: Comparisons to state-of-the-art structured variational autoencoder baselines show improvements in terms of the expressiveness of the learned model.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Deep latent variable models offer an effective method for automatically learning structure from data.
Revised: Deep latent variable models offer an effective method for automatically learning structure from data.

Original: By explicitly modeling the data distribution using latent variables, these models are capable of learning compressed representations that are then relevant for downstream tasks.
Revised: By explicitly modeling the data distribution using latent variables, these models are capable of learning compressed representations that are then relevant for downstream tasks.

Original: Such models have been applied across a wide array of domains, such as images (Gregor et al., 2014; Kingma &
Revised: Such models have been applied across a wide array of domains, such as images (Gregor et al., 2014; Kingma &

Original: Welling, 2014; Rezende et al., 2014; Gregor et al., 2015), audio (Chung et al., 2015; Fraccaro et al., 2016), video (He et al., 2018;
Revised: Welling, 2014; Rezende et al., 2014; Gregor et al., 2015), audio (Chung et al., 2015; Fraccaro et al., 2016), video (He et al., 2018;

Original: Yingzhen & Mandt, 2018), and text (Bowman et al., 2016; Krishnan et al., 2017).
Revised: Yingzhen & Mandt, 2018), and text (Bowman et al., 2016; Krishnan et al., 2017).

Original: However, despite their success, latent variable models are often formulated with simple (e.g. Gaussian) distributions, making independence assumptions among the latent variables.
Revised: However, despite their success, latent variable models are often formulated with simple (e.g. Gaussian) distributions, making independence assumptions among the latent variables.

Original: That is, each latent variable is sampled independently.
Revised: That is, each latent variable is sampled independently.

Original: Ignoring the dependencies between latent variables limits the ﬂexibility of these models, negatively impacting the model’s ability to ﬁt the data.
Revised: Ignoring the dependencies between latent variables limits the ﬂexibility of these models, negatively impacting the model’s ability to ﬁt the data.

Original: Latent dependencies can be incorporated both during inference as well as in the model itself.
Revised: Latent dependencies can be incorporated both during inference as well as in the model itself.

Original: For instance, normalizing ﬂows (Rezende & Mohamed, 2015) accounts for latent dependencies during inference by learning a mapping from a simple distribution to a more complex distribution that contains these dependencies.
Revised: For instance, normalizing ﬂows (Rezende & Mohamed, 2015) accounts for latent dependencies during inference by learning a mapping from a simple distribution to a more complex distribution that contains these dependencies.

Original: On the modeling side, one can directly add dependencies by constructing a hierarchical latent representation (Dayan et al., 1995).
Revised: On the modeling side, one can directly add dependencies by constructing a hierarchical latent representation (Dayan et al., 1995).

Original: These structures consist of conditional ( empirical ) priors, in which one latent variable forms a prior on another latent variable.
Revised: These structures consist of conditional ( empirical ) priors, in which one latent variable forms a prior on another latent variable.

Original: While this conditional distribution may take a simple form, marginalizing over the parent variable can result in an arbitrarily complex distribution.
Revised: While this conditional distribution may take a simple form, marginalizing over the parent variable can result in an arbitrarily complex distribution.

Original: Models with these more ﬂexible latent dependency structures have been shown to result in improved performance (Sønderby et al., 2016; Burda et al., 2016;
Revised: Models with these more ﬂexible latent dependency structures have been shown to result in improved performance (Sønderby et al., 2016; Burda et al., 2016;

Original: Kingma et al., 2016).
Revised: Kingma et al., 2016).

Original: However, despite the beneﬁts of including additional structure in these models, latent dependency structures have been predeﬁned, potentially limiting the performance of this approach.
Revised: However, despite the beneﬁts of including additional structure in these models, latent dependency structures have been predeﬁned, potentially limiting the performance of this approach.

Original: In this work, we propose a method for learning dependency structures in latent variable models.
Revised: In this work, we propose a method for learning dependency structures in latent variable models.

Original: Structure learning is a difﬁcult task with a long history in the graphical models community (Koller & Friedman, 2009).
Revised: Structure learning is a difﬁcult task with a long history in the graphical models community (Koller & Friedman, 2009).

Original: Over the years, it has been tackled from several perspectives, including constraint-based approaches (Cheng et al., 2002; Lehmann & Romano, 2008), optimization of structure scores (Kass & Raftery, 1995; Heckerman et al., 1995; Barron et al., 1998), Bayesian model averaging (Heckerman et al., 1999; Koivisto & Sood, 2004), and many more.
Revised: Over the years, it has been tackled from several perspectives, including constraint-based approaches (Cheng et al., 2002; Lehmann & Romano, 2008), optimization of structure scores (Kass & Raftery, 1995; Heckerman et al., 1995; Barron et al., 1998), Bayesian model averaging (Heckerman et al., 1999; Koivisto & Sood, 2004), and many more.

Original: Unfortunately, the underlying objectives are often limited to graphs of a particular form (e.g., limited tree width), prohibitively expensive, or difﬁcult to integrate with the gradient-based optimization techniques of modern neural networks.
Revised: Unfortunately, the underlying objectives are often limited to graphs of a particular form (e.g., limited tree width), prohibitively expensive, or difﬁcult to integrate with the gradient-based optimization techniques of modern neural networks.

Original: Here, we discuss a variational end-to-end approach for general graph structures introducing minimal complexity overhead.
Revised: Here, we discuss a variational end-to-end approach for general graph structures introducing minimal complexity overhead.

Original: In particular, we introduce a set of binary global variables to gate the latent dependencies.
Revised: In particular, we introduce a set of binary global variables to gate the latent dependencies.

Original: The whole model (including its structure) is jointly optimized with a single stochastic variational inference objective.
Revised: The whole model (including its structure) is jointly optimized with a single stochastic variational inference objective.

Original: In our experimental validation, we show that the learned dependency structures contribute to a more accurate representation of the true generative distribution, outperforming several other variants of variational autoencoders.
Revised: In our experimental validation, we show that the learned dependency structures result in models that more accurately model the data distribution, outperforming several common predeﬁned latent dependency structures.

Original:  B ACKGROUND 2. V ARIATIONAL I NFERENCE & V ARIATIONAL A UTOENCODERS
Revised:  B ACKGROUND 2. V ARIATIONAL I NFERENCE & V ARIATIONAL A UTOENCODERS

Original: A latent variable model, deﬁned by the joint distribution, p θ ( x , z ) =
Revised: A latent variable model, deﬁned by the joint distribution, p θ ( x , z ) =

Original: p θ ( x ∣ z ) p θ ( z ) , models each data example, x , using a local latent variable, z , and global parameters, θ .
Revised: p θ ( x ∣ z ) p θ ( z ) , models each data example, x , using a local latent variable, z , and global parameters, θ .

Original: p θ ( x ∣ z ) denotes the conditional likelihood, and p θ ( z ) denotes the prior.
Revised: p θ ( x ∣ z ) denotes the conditional likelihood, and p θ ( z ) denotes the prior.

Original: Latent variable models are capable of capturing the structure present in data, with z forming a compressed representation of each data example.
Revised: Latent variable models are capable of capturing the structure present in data, with z forming a compressed representation of each data example.

Original: Unfortunately, inferring the posterior, p θ ( z ∣ x ) , is typically computationally intractable, prompting the use of approximate inference techniques.
Revised: Unfortunately, inferring the posterior, p θ ( z ∣ x ) , is typically computationally intractable, prompting the use of approximate inference techniques.

Original: Variational inference (Jordan et al., 1999) introduces an approximate posterior, q φ (
Revised: Variational inference (Jordan et al., 1999) introduces an approximate posterior, q φ (

Original: z ∣ x ) , and optimizes variational parameters, φ , to minimize the KL-divergence to the true posterior, KL ( q φ (
Revised: z ∣ x ) , and optimizes variational parameters, φ , to minimize the KL-divergence to the true posterior, KL ( q φ (

Original: z ∣ x )∣∣
Revised: z ∣ x )∣∣

Original: p θ ( z ∣ x )) .
Revised: p θ ( z ∣ x )) .

Original: As this quantity cannot be evaluated directly, the following relation is used:
Revised: As this quantity cannot be evaluated directly, the following relation is used:

Original: [Equation]
Revised: [Equation]

Original: where L( x ; θ,φ ) is the evidence lower bound (ELBO), deﬁned as
Revised: where L( x ; θ,φ ) is the evidence lower bound (ELBO), deﬁned as

Original: [Equation]
Revised: [Equation]

Original: In Eq.
Revised: In Eq.

Original: (1), log p θ ( x ) is independent of φ , so we can minimize the KL divergence term, i.e. perform approximate inference, by maximizing L( x ; θ,φ ) w.r.t.
Revised: (1), log p θ ( x ) is independent of φ , so we can minimize the KL divergence term, i.e. perform approximate inference, by maximizing L( x ; θ,φ ) w.r.t.

Original: q φ ( z ∣ x ) .
Revised: q φ ( z ∣ x ) .

Original: Further, because KL divergence is non-negative, L( x ; θ,φ ) is a lower bound on log p θ ( x ) , meaning we can then learn the model parameters by maximizing L( x ; θ,φ ) w.r.t.
Revised: Further, because KL divergence is non-negative, L( x ; θ,φ ) is a lower bound on log p θ ( x ) , meaning we can then learn the model parameters by maximizing L( x ; θ,φ ) w.r.t.

Original: θ .
Revised: θ .

Original: Variational autoencoders (VAEs) (Kingma & Welling, 2014; Rezende et al., 2014) amortize inference optimization across data examples by parameterizing q φ ( z ∣ x ) as a separate inference model, then jointly optimizing the model parameters θ and φ .
Revised: Variational autoencoders (VAEs) (Kingma & Welling, 2014; Rezende et al., 2014) amortize inference optimization across data examples by parameterizing q φ ( z ∣ x ) as a separate inference model, then jointly optimizing the model parameters θ and φ .

Original: VAEs instantiate both the inference model and latent variable model with deep networks, allowing them to scale to high-dimensional data.
Revised: VAEs instantiate both the inference model and latent variable model with deep networks, allowing them to scale to high-dimensional data.

Original: However, VAEs are typically implemented with basic graphical structures and simple, unimodal distributions (e.g. Gaussians).
Revised: However, VAEs are typically implemented with basic graphical structures and simple, unimodal distributions (e.g. Gaussians).

Original: For instance, the dimensions of the prior are often assumed to be independent, p θ ( z )
Revised: For instance, the dimensions of the prior are often assumed to be independent, p θ ( z )

Original: =
Revised: =

Original: ∏ m p
Revised: ∏ m p

Original: θ
Revised: θ

Original: z m ) , with a common assumption being a ﬁxed standard Gaussian: p θ ( z ) =
Revised: z m ) , with a common assumption being a ﬁxed standard Gaussian: p θ ( z ) =

Original: N( z ; 0 , I ) .
Revised: N( z ; 0 , I ) .

Original: Similarly, approximate posteriors often make the mean ﬁeld assumption, q φ ( z ∣ x )
Revised: Similarly, approximate posteriors often make the mean ﬁeld assumption, q φ ( z ∣ x )

Original: =
Revised: =

Original: ∏ m q φ ( z m ∣ x ) .
Revised: ∏ m q φ ( z m ∣ x ) .

Original: Independence assumptions such as these may be overly restrictive, thereby limiting modeling capabilities.
Revised: Independence assumptions such as these may be overly restrictive, thereby limiting modeling capabilities.

Original: 2. E MPIRICAL P RIORS T HROUGH L ATENT D EPENDENCY S TRUCTURE
Revised: 2. E MPIRICAL P RIORS T HROUGH L ATENT D EPENDENCY S TRUCTURE

Original: One technique for improving the expressive capacity of latent variable models is through incorporating dependency structure among the latent variables, forming a hierarchy (Dayan et al., 1995;
Revised: One technique for improving the expressive capacity of latent variable models is through incorporating dependency structure among the latent variables, forming a hierarchy (Dayan et al., 1995;

Original: Rezende et al., 2014).
Revised: Rezende et al., 2014; Goyal et al., 2017; Vikram et al., 2018; Webb et al., 2018).

Original: These dependencies provide empirical priors, learned priors that are conditioned on other latent variables.
Revised: These dependencies provide empirical priors, learned priors that are conditioned on other latent variables.

Original: With M latent dimensions, the full prior takes the following auto-regressive form:
Revised: With M latent dimensions, the full prior takes the following auto-regressive form:

Original: [Equation]
Revised: [Equation]

Original: ) where z pa ( m ) denotes the vector of latent variables constituting the parents of z m .
Revised: where z pa ( m ) denotes the vector of latent variables constituting the parents of z m .

Original: Each conditional distribution can be parameterized by deep networks that output the parameters of distributions, e.g. mean and variance of a Gaussian.
Revised: Each conditional distribution can be parameterized by deep networks that output the parameters of distributions, e.g. mean and variance of a Gaussian.

Original: While these conditional distributions may be relatively simple, the marginal empirical prior, p θ ( z m ) = ∫ p θ ( z m ∣ z pa ( m ) )
Revised: While these conditional distributions may be relatively simple, the marginal empirical prior , p θ ( z m ) = ∫ p θ [Equation] z m ∣ z pa ( m ) )

Original: p θ ( z pa ( m ) )
Revised: p θ ( z pa ( m ) )

Original: d z pa ( m ) , can be arbitrarily complex.
Revised: d z pa ( m ) , can be arbitrarily complex.

Original: By using more ﬂexible priors, models with latent dependency structure are less restrictive in their latent representations, hopefully capturing more information and enabling a better ﬁt to the data.
Revised: By using more ﬂexible priors, models with latent dependency structure are less restrictive in their latent representations, hopefully capturing more information and enabling a better ﬁt to the data.

Original: [Table]
Revised: [Table]

Original: z )
Revised: z )

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Figure]
Revised: [Figure]

Original: With the added latent dependencies in the model, the independence assumption in the approximate posterior is even less valid in this setting.
Revised: With the added latent dependencies in the model, the independence assumption in the approximate posterior is even less valid in this setting.

Original: While normalizing ﬂows (Rezende & Mohamed, 2015) offers one technique for overcoming the mean ﬁeld assumption, a separate line of work has investigated the use of structured approximate posteriors, particularly in the context of models with empirical priors (Johnson et al., 2016).
Revised: While normalizing ﬂows (Rezende & Mohamed, 2015) offers one technique for overcoming the mean ﬁeld assumption, a separate line of work has investigated the use of structured approximate posteriors, particularly in the context of models with empirical priors (Johnson et al., 2016).

Original: This technique introduces dependencies between the dimensions of the approximate posterior, often mirroring the dependency structure of the latent variable model.
Revised: This technique introduces dependencies between the dimensions of the approximate posterior, often mirroring the dependency structure of the latent variable model.

Original: An explanation for this was provided by Marino et al.
Revised: An explanation for this was provided by Marino et al.

Original: (2018): optimizing the approximate posterior requires knowledge of the prior, which is especially relevant in models with empirical priors where the prior can vary with the data.
Revised: (2018): optimizing the approximate posterior requires knowledge of the prior, which is especially relevant in models with empirical priors where the prior can vary with the data.

Original: Ladder VAE (Sønderby et al., 2016) incorporates these prior dependencies by using a structured approximate posterior of the form
Revised: Ladder VAE (Sønderby et al., 2016) incorporates these prior dependencies by using a structured approximate posterior of the form

Original: [Equation]
Revised: [Equation]

Original: Unlike the mean ﬁeld approximate posterior, which conditions each dimension only on the data example, x , the distributions in Eq.
Revised: Unlike the mean ﬁeld approximate posterior, which conditions each dimension only on the data example, x , the distributions in Eq.

Original: (4) account for latent dependencies by conditioning on samples from the parent variables.
Revised: (4) account for latent dependencies by conditioning on samples from the parent variables.

Original: Ladder VAE performs this conditioning by reusing the empirical prior during inference, forming the approximate posterior by combining a “bottom-up” recognition distribution and the “top-down” prior.
Revised: Ladder VAE performs this conditioning by reusing the empirical prior during inference, forming the approximate posterior by combining a “bottom-up” recognition distribution and the “top-down” prior.

Original: While Eqs.
Revised: While Eqs.

Original: (3) and (4) permit separate latent dependencies for each individual latent dimension, the dimensions are typically partitioned into a set of nodes, with dimensions within each node sharing the same parents.
Revised: (3) and (4) permit separate latent dependencies for each individual latent dimension, the dimensions are typically partitioned into a set of nodes, with dimensions within each node sharing the same parents.

Original: This improves computational efﬁciency by allowing priors and approximate posteriors within each node to be calculated in parallel.
Revised: This improves computational efﬁciency by allowing priors and approximate posteriors within each node to be calculated in parallel.

Original: Using z n to denote latent node n of N and z pa ( n ) to denote the concatenation of its parent nodes, we can write the ELBO (Eq.
Revised: Using z n to denote latent node n of N and z pa ( n ) to denote the concatenation of its parent nodes, we can write the ELBO (Eq.

Original: (2)) as
Revised: (2)) as

Original: [Equation]
Revised: [Equation]

Original: Note that the KL divergence term in the ELBO can no longer be evaluated analytically, now requiring a sampling-based estimate of the expectation (Kingma & Welling, 2014).
Revised: Note that the KL divergence term in the ELBO can no longer be evaluated analytically, now requiring a sampling-based estimate of the expectation (Kingma & Welling, 2014).

Original: While this can lead to higher variance in ELBO estimates and the resulting gradients, models with latent dependencies still tend to empirically outperform models with independence assumptions (Burda et al., 2016;
Revised: While this can lead to higher variance in ELBO estimates and the resulting gradients, models with latent dependencies still tend to empirically outperform models with independence assumptions (Burda et al., 2016;

Original: Sønderby et al., 2016).
Revised: Sønderby et al., 2016).

Original: However, by increasing the number of nodes, the burden of devising a suitable dependency structure falls upon the experimental practitioner.
Revised: However, by increasing the number of nodes, the burden of devising a suitable dependency structure falls upon the experimental practitioner.

Original: This is non-trivial, as the structure may depend on the data and other model hyperparameters, such as the number of layers in the deep networks, non-linearities, latent distributions, etc.
Revised: This is non-trivial, as the structure may depend on the data and other model hyperparameters, such as the number of layers in the deep networks, non-linearities, latent distributions, etc.

Original: Rather than relying on pre-deﬁned fully-connected structures (Kingma et al., 2016) or chain structures (Sønderby et al., 2016), we seek to automatically learn the latent dependency structure as part of the variational optimization process.
Revised: Rather than relying on pre-deﬁned fully-connected structures (Kingma et al., 2016) or chain structures (Sønderby et al., 2016), we seek to automatically learn the latent dependency structure as part of the variational optimization process.

Original: A comparison of these approaches is visualized in Fig.
Revised: A comparison of these approaches is visualized in Fig.

Original:  V ARIATIONAL O PTIMIZATION OF L ATENT S TRUCTURES
Revised:  V ARIATIONAL O PTIMIZATION OF L ATENT S TRUCTURES

Original: A major difﬁculty in the context of variational optimization of latent structures is the fact that, unlike the model parameters ( φ,θ ) , which are optimized over a continuous domain, the latent dependency structure is discrete , without a clear ordering.
Revised: Unlike the model parameters ( φ,θ ) , which are optimized over a continuous domain, the latent dependency structure is discrete , without a clear ordering.

Original: While the discrete nature of the latent space’s topological structure introduces discontinuities in the optimization landscape, complicating the learning process, there is (unlike the related setting of neural architecture search (Zoph & Le, 2016)) fortunately only a ﬁnite number of possible dependency structures over a ﬁxed number of latent dimensions:
Revised: The discrete nature of the latent space’s topological structure introduces discontinuities in the optimization landscape, complicating the learning process. Fortunately, unlike the related setting of neural architecture search (Zoph & Le, 2016), there is only a ﬁnite number of possible dependency structures over a ﬁxed number of latent dimensions:

Original: In a directed graphical model, a fully-connected directed acyclic graph (DAG) models all possible dependencies.
Revised: In a directed graphical model, a fully-connected directed acyclic graph (DAG) models all possible dependencies.

Original: In this model, an ordering is induced over the latent nodes, and the parents of node n (of N ) are given as z pa ( n ) =
Revised: In this model, an ordering is induced over the latent nodes, and the parents of node n (of N ) are given as z pa ( n ) =

Original: { z n + 1 ,..., z N } .
Revised: { z n + 1 ,..., z N } .

Original: Thus, to learn a latent dependency structure, we can maintain all dependencies in a fully-connected DAG, modifying their presence or absence during training.
Revised: Thus, to learn an appropriate latent dependency structure, we can maintain all dependencies in a fully-connected DAG, modifying their presence or absence during training.

Revised: This is accomplished by introducing a set of binary dependency gates (Section 3.1).

Original: Naively, one might assume that a fully-connected DAG is the most general structure, and, therefore, current models should implicitly learn to ignore unnecessary dependencies. However, latent variable models are highly prone to local optima (Bowman et al., 2016; Burda et al., 2016), and as we show empirically in Section 4, modifying entire dependencies can yield improved performance.
Revised: We convert discrete optimization over dependency structures into a continuous optimization problem by parameterizing these gates as samples from Bernoulli distributions, then learning the distribution parameters (Section 3.3). These gating distributions induce an additional lower bound on L , which becomes tight when the distribution converges to a delta function, yielding a single, optimized dependency structure (Section 3.2). Indeed, we observe this process empirically, with the learned dependency structures outperforming their predeﬁned counterparts (Section 4).

Original: 3. G ATED D EPENDENCIES
Revised: 3. G ATED D EPENDENCIES

Original: To control the dependency structure of the model, we introduce a set of binary global variables, c
Revised: To control the dependency structure of the model, we introduce a set of binary global variables, c

Original: = { c i,j } i,j , which gate the latent dependencies.
Revised: = { c i,j } i,j , which gate the latent dependencies.

Original: The element c i,j denotes the gate variable from z i to z j
Revised: The element c i,j denotes the gate variable from z i to z j

Original: ( i > j ), specifying the presence or absence of this latent dependency.
Revised: ( i > j ), specifying the presence or absence of this latent dependency.

Original: We treat each c i,j as an independent random variable, sampled from a Bernoulli distribution with mean µ i,j ,

Original: i.e. c i,j ∼ B( µ i,j ) .

Original: We denote the set of these Bernoulli means as µ .

Original: Because each element of c takes values in { 0 , 1 } , dependencies can be preserved or removed simply through (broadcasted) element-wise multiplication with the corresponding parent nodes.
Revised: Because each element of c takes values in { 0 , 1 } , dependencies can be preserved or removed simply through (broadcasted) element-wise multiplication with the corresponding parent nodes.

Original: Removing a dependency entails multiplying the corresponding input parent node by 0 .
Revised: Removing a dependency entails multiplying the corresponding input parent node by 0 .

Original: With this addition, the model is now  expressed as p θ ( x , z , c ) =
Revised: Each possible latent dependency structure can now be expressed through its corresponding value of c .

Original:  p θ ( x ∣ z , c )  p θ ( z ∣  c ) p  (  c ) . Whereas the model parameters, θ ,are point estimates, the dependency gates are governed by Bernoulli distributions, with each sampling from p ( c ) inducing a corresponding dependency structure.
Revised: Each ﬁxed latent dependency structure, c ′ , induces a separate latent variable model p θ ( x , z , c ) = p θ ( x ∣ z , c ) p θ ( z ∣ c ) δ c , c ′ , where δ ⋅ , ⋅ is the Kronecker delta, which effectively selects a single structure.

Original: Similar to Eq.
Revised: Similar to Eq.

Original: (3), the prior  can now be expressed as
Revised: (3), the prior on the latent variables can now be expressed as

Original: [Equation]
Revised: [Equation]

Original: (6) where c pa ( n ) ,n denotes the gate variables associated with the dependencies between node z n and its parents, z pa ( n ) .
Revised: (6) where c pa ( n ) ,n denotes the gate variables associated with the dependencies between node z n and its parents, z pa ( n ) .

Original: Note that z pa ( n ) denotes the set of all possible parents of node z n in the fullyconnected DAG, i.e. z pa ( n ) =
Revised: Note that z pa ( n ) denotes the set of all possible parents of node z n in the fullyconnected DAG, i.e. z pa ( n ) =

Original: { z n + 1 ,..., z N } .
Revised: { z n + 1 ,..., z N } .

Original: To give a concrete example of the gating procedure, consider the case in which p θ ( z n ∣ z pa ( n ) , c pa ( n ) ,n ) is given by a Gaussian density with parameters
Revised: To give a concrete example of the gating procedure, consider the case in which p θ ( z n ∣ z pa ( n ) , c pa ( n ) ,n ) is given by a Gaussian density with parameters

Original: ̂ ψ n = (̂ µ n , ̂ Σ n ) .
Revised: ̂ ψ n = (̂ µ n , ̂ Σ n ) .

Original: We obtain these parameters recursively by multiplying samples of node z n ’s parent variables z pa(n) withsamples of their corresponding gating variables c pa(n),n and inject a concatentation of the results of this operation into a multi-layer perceptron MLP ( TD ) n predicting ̂ ψ n (see 
Revised: We obtain these parameters recursively by multiplying samples of node z n ’s parent variables z pa(n) with their corresponding gating variables c pa(n),n and input a concatenation of the results of this operation into a multi-layer perceptron MLP ( TD ) n predicting ̂ ψ n (see Appendix

Original: Appendix A.2 for additional details on the MLP architecture).
Revised: B for additional details on the MLP architecture).

Original: The top-down recursion ends at the root node z N ∼  N( 0 , I ) .
Revised: The top-down recursion starts at the root node z N ∼ p ( z N ) = N( 0 , I ) .

Original: An illustration of this process is shown in black in Fig.
Revised: An illustration of this process is shown in black in Fig.

Original: The approximate posterior, q φ ( z ∣ x  ) , must approximate p θ ( z ∣ x) ,which is marginalized across c  . 
Revised: The approximate posterior, q φ ( z ∣ x , c ) , must approximate p θ ( z ∣ x , c ) . We express the approximate posterior as

Original: In effect, q φ ( z ∣ x ) should account for the entire distribution of dependency structures, rather than any
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: F <latexit
Revised: F <latexit

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: -conditioned parameters ̃ ψ n from a generic encoding of x .particular structure.
Revised: -conditioned parameters ̃ ψ n from a generic encoding of x .

Original: Thus, we express the approximate posterior as: [Equation] where z pa ( n ) denotes the parent nodes, marginalized across the dependency gate variables, c .
Revised: While Webb et al.

Original: In practice, this involves sampling dependency gating variables to estimate the top-down input for inference. With a single sample, corresponding to a single dependency structure, this exactly corresponds to  the inference procedure from Sønderby et al.
Revised: (2018) provide the form of the approximate posterior dependencies implied by a given latent dependency structure, we restrict the approximate posterior to mirror the dependencies in the generative model. This reduces the number of learned parameters, and we found that it gave improved empirical performance over maintaining all dependencies.

Original: Parameter prediction for the local factors q φ ( z n ∣ x , z pa ( n )  ) consists of a precision-weighted fusion of the top-down prediction ̂
Revised: Parameter prediction for the local factors q φ ( z n ∣ x , z pa ( n ) , c pa ( n ) ) consists of a precision-weighted fusion of the top-down prediction ̂

Original: ψ n described above and a bottom-up prediction ̃ ψ n .
Revised: ψ n described above and a bottom-up prediction ̃ ψ n .

Original: The latter is obtained by encoding x into a generic feature that is used as an input to a node-speciﬁc multi-layer perceptron MLP ( BU ) n predicting ̃ ψ n .
Revised: The latter is obtained by encoding x into a generic feature that is used as an input to a node-speciﬁc multi-layer perceptron MLP ( BU ) n predicting ̃ ψ n .

Original: This is shown in blue in Fig.
Revised: This is shown in blue in Fig.

Original: Introducing the  dependency gating variables modiﬁes the variational objective (Eq.
Revised: Additional details on the fusion process can be found in Appendix B.3. 3. L EARNING S TRUCTURE BY I NDUCING AN A DDITIONAL L OWER B OUND The formulation in Section 3.1 provides the form of the model for a particular conﬁguration of the latent dependency structure. Finding the optimal structure corresponds to a discrete optimization over all values of c , potentially optimizing the model parameters of each possible conﬁguration.

Original: (2)), as we must now marginalize over these additional variables.
Revised: To avoid this intractable procedure, we place a distribution over c , then directly optimize the parameters of this distribution to arrive at a single, learned latent dependency structure. Speciﬁcally, we treat each c i,j as an independent random variable, sampled from a Bernoulli distribution with mean µ i,j , i.e. c i,j ∼ p ( c i,j ) = B( µ i,j ) .

Revised: We denote the set of these Bernoulli means as µ .

Original: In Appendix A.1, we derive the corresponding lower bound, ̃L , which can be expressed as [Equation]
Revised: Introducing this distribution allows us to express the following additional lower bound on L , derived in Appendix A:

Original: [Equation]
Revised: [Equation]

Original: where L c is the ELBO for a given value of dependency gating variables, c .
Revised: where L c is the ELBO for a particular value of dependency gating variables.

Original: Thus, ̃L can be interpreted as the expected ELBO under the distribution of dependency structures induced by p ( c )  . 
Revised: Thus, ̃L can be interpreted as the expected ELBO under the distribution of dependency structures induced by p ( c ) , which we estimate by sampling c ∼ p ( c ) and evaluating L c . We note that ̃L is not a proper variational bound, as it is not guaranteed to recover the marginal likelihood if the approximate posterior matches the true posterior. Rather, optimizing ̃L provides a method for learning the latent structure. For ﬁxed parameters ( φ,θ ) , the optimal ̃ L w.r.t.

Original: We form a Monte Carlo estimate of ̃L by sampling c ∼ p ( c ) and evaluating L c  .
Revised: µ is a δ -distribution at the MAP conﬁguration, c ∗ , yielding L = ̃L = L c ∗ .

Revised: This is because L c ∗ is always greater than or equal to the expected ELBO over all dependency gates, ̃L .

Original:  The training procedure is outlined in Algorithm 1.
Revised: In practice, we jointly optimize φ , θ , and µ . While this non-convex optimization procedure may result in local optima, we ﬁnd this empirically works well, with p ( c ) converging to a ﬁxed distribution (Fig. 3). Thus, by the end of training, we are effectively optimizing the ELBO for a single dependency structure. The training procedure is outlined in Algorithm 1.

Original: 3. L EARNING THE D ISCRETE G ATING V ARIABLE D ISTRIBUTIONS
Revised: 3. L EARNING THE D ISCRETE G ATING V ARIABLE D ISTRIBUTIONS

Original: For a given latent dependency structure, gradients for the parameters θ and φ can be estimated using
Revised: For a given latent dependency structure, gradients for the parameters θ and φ can be estimated using

Original: Monte Carlo samples and the reparameterization trick (Kingma & Welling, 2014; Rezende et al., 2014).
Revised: Monte Carlo samples and the reparameterization trick (Kingma & Welling, 2014; Rezende et al., 2014).

Original: To obtain gradients for the gate means, µ , we make use of recent advances in differentiating through discrete operations (Maddison et al., 2017; Jang et al., 2017), allowing us to differentiate through the sampling of the dependency gating variables, c .
Revised: To obtain gradients for the gate means, µ , we make use of recent advances in differentiating through discrete operations (Maddison et al., 2017; Jang et al., 2017), allowing us to differentiate through the sampling of the dependency gating variables, c .

Original: Speciﬁcally, we recast the gating variables using the Gumbel-Softmax estimator from Jang et al.
Revised: Speciﬁcally, we recast the gating variables using the Gumbel-Softmax estimator from Jang et al.

Original: (2017), re-expressing c
Revised: (2017), re-expressing c

Original: i,j as:
Revised: i,j as:

Original: [Equation]
Revised: [Equation]

Original: Data latent nodes where (cid:15) 1,(cid:15) 2 , and (cid:15) 3 are i.i.d samples drawn from a Gumbel( 0 , 1 ) distribution and τ is a temperature parameter.
Revised: (9) where (cid:15) 1 and (cid:15) 2 are i.i.d samples drawn from a Gumbel( 0 , 1 ) distribution and τ is a temperature parameter.

Original: The Gumbel-Softmax distribution is differentiable for τ > 0 , allowing us to estimate the derivative
Revised: The Gumbel-Softmax distribution is differentiable for τ > 0 , allowing us to estimate the derivative

Original: ∂c i,j ∂µ i,j .
Revised: ∂c i,j ∂µ i,j .

Original: For large values of τ , we obtain smoothed versions of c , essentially interpolating between different dependency structures.
Revised: For large values of τ , we obtain smoothed versions of c , essentially interpolating between different dependency structures.

Original: As τ → 0 , we recover binary values for c , yielding the desired discrete sampling of dependency structures at the cost of high-variance gradient estimates.
Revised: As τ → 0 , we recover binary values for c , yielding the desired discrete sampling of dependency structures at the cost of high-variance gradient estimates.

Original: Thus, we anneal τ during training to learn the dependency gate means, µ , eventually arriving at the discrete setting.
Revised: Thus, we anneal τ during training to learn the dependency gate means, µ , eventually arriving at the discrete setting.

Original:  E VALUATION
Revised:  E VALUATION

Original: We evaluate the proposedmodel and its learned latent dependency structure on three challenging datasets: MNIST (Lecun et al., 1998; Larochelle & Murray, 2011), Omniglot (Lake et al., 2013), and CIFAR-10 (Krizhevsky, 2009).
Revised: We evaluate the proposed latent dependency learning approach on three benchmark datasets: MNIST (Lecun et al., 1998; Larochelle & Murray, 2011), Omniglot (Lake et al., 2013), and CIFAR10 (Krizhevsky, 2009).

Original: MNIST and Omniglot are binarized and modeled with Bernoulli distributions; CIFAR-10 is modeled with Gaussian densities. Our baselines consist of classic variational autoencoders (Kingma & Welling, 2014; Rezende et al., 2014) as well as its popular variants, including ladder VAEs (Sønderby et al., 2016) and VAEs with fully-connected latent dependency structure (FC-VAEs; cf. (Kingma et al., 2016)). Together, these baselines cover the full spectrum of latent dependency structures, ranging from the strongest possible independence assumptions (fully independent in classic VAEs) to linear dependencies (chainstructured in ladder VAEs) to  the most expressive representation (fully connected in FC-VAEs).
Revised: After discussing the experimental setup in Section 4.1, we provide a set of qualitative experiments in Section 4.2 to gain insight into the learning process, hyper-parameter selection, and the nature of the inferred structures.

Original: Our experiments show that none of those predeﬁned extremes is optimal. Instead, we demonstrate that learning a data-driven distribution over graph structures that is jointly optimized with the variational and generative model parameters leads to a superior representation of the underlying test-time distribution.
Revised: In Section 4.3, we provide quantitative comparisons with common predeﬁned latent dependency structures on benchmark datasets.

Original: An ablation study provides further insights into our learning process, hyperparameter selection, and the nature of the inferred structures. We note that our approach is orthogonal andcould be complemented by a number of other approaches attempting to overcome the mean ﬁeld assumptions of classic VAEs: Similar to ladder VAEs, (Zhao et al., 2017) also introduces chain-structured latent dependencies but focuses on learning a disentangled representation. VAEs with normalizing ﬂows (Rezende & Mohamed, 2015) gradually increase the complexity of the approximate posterior through a series of invertible transformations.
Revised: Additional results on training robustness and latent space embeddings can be found in Appendix C.

Original: 4. E XPERIMENTAL S ETUP
Revised: 4. E XPERIMENTAL S ETUP

Original: To guarantee a fair comparison, the encoders of all structured methods use the same MLP architecture with batch normalization (Ioffe & Szegedy, 2015) and ReLU non-linearities (Nair & Hinton, 2010) in all experiments.
Revised: To provide a fair comparison, the encoders of all structured methods use the same MLP architecture with batch normalization (Ioffe & Szegedy, 2015) and ReLU non-linearities (Nair & Hinton, 2010) in all experiments.

Original: 2 Decoder structures are the reverse of the encoders.
Revised: 2 Decoder structures are the reverse of the encoders.

Original: Likewise, the number of latent dimensions is the same in all experiments/models ( M = 80 ).
Revised: Likewise, the number of latent dimensions is the same in all models and experiments ( M = 80 ).

Original: As discussed in section 3.1, all latent dependencies are modelled by non-linear MLPs as well.
Revised: As discussed in Section 3.1, all latent dependencies are modeled by non-linear MLPs as well.

Original: For MNIST and Omniglot, we  model p θ ( x ∣ z ) as a Bernoulli distribution, using a sigmoid nonlinearity on the output layer to produce the mean of this distribution.
Revised: For MNIST and Omniglot, we binarize the data and model p θ ( x ∣ z ) as a Bernoulli distribution, using a sigmoid non-linearity on the output layer to produce the mean of this distribution.

Original: For CIFAR-10, we model 
Revised: For CIFAR-10, we model p θ ( x ∣ z ) with a Gaussian density, with mean and log-variance predicted by sigmoid and linear functions, respectively. Further implementation details, including the model architectures and training criteria, can be found in Appendix B. Accompanying code will be made publicly available after the anonymous review period.

Original: [Figure]
Revised: [Figure]

Original: 4. Q UALITATIVE A NALYSIS
Revised: 4. Q UALITATIVE A NALYSIS

Revised: We ﬁrst explore the structure learning process.

Revised: As described in Section 3.2 and Appendix A, optimizing ̃L w.r.t.

Revised: the dependency gating means, µ , should push this lower bound toward L .

Revised: Thus, µ should converge to either 0 or 1 , yielding a ﬁxed, ﬁnal latent dependency structure.

Original: An example of the structure learning process  on MNIST with N = 5 nodes and a total latent dimension of M = 

Revised: In Fig. 3, we visualize this process during training on MNIST. The model has N = 5 nodes and a total latent dimension of M =

Original: (i e., N ′ =
Revised: (i e., N ′ =

Original: M / N = 16 dimensions per node) is visualized in Fig.
Revised: M / N = 16 dimensions per node).

Original: In Fig.
Revised: As shown in Fig.

Original: 3a,we show the evolution of the parameters µ of the gating variables c . The network actively drops 3 /edges during the learning process, while the remaining parameters eventually converge to 1 .
Revised: 3a, the gating means converge in practice with 3 out of 10 edges removed and the rest retained.

Original: We observed that the resulting distribution is stable across multiple training runs with different seeds for weight initialization and minibatch sampling.

Original: A sample of the learned distribution is shown in Fig.
Revised: The resulting static dependency structure is visualized in Fig.

Original: 3b
Revised: 3b.

Original: Fig. 4 reports an ablation study conducted on MNIST.
Revised: We observed that the learned structure is stable across training runs with different seeds for parameter initialization and mini-batch sampling, suggesting that it depends on the model parameterization and the dataset.

Original: In particular, we investigate the inﬂuence of the total latentdimension M and, for ﬁxed latent dimension,  the trade-off between  number of nodes N and  node dimension N ′ =
Revised: We next investigate the inﬂuence of the total latent dimension, M , and the trade-off between the number of nodes, N and the node dimension, N ′ =

Original: M / N .
Revised: M / N .

Revised: Results are shown in Fig.

Revised: 4 for various models trained on MNIST.

Original: Models with the same total latent dimensionM are shown in the same color.
Revised: Models with the same total latent dimension are shown in the same color.

Original: We observe two trends: (1) The performance of our model increases with latent dimension.

Original: This is an expected result and a likely explanation for this behaviour is the additional ﬂexibility of a higher-dimensional latent space; (2) For a ﬁxed number of latent dimensions, less node dimensions (and thus more nodes and the potential for a more complex  dependency structure) typically performs better, highlighting the importance of a learned, non-trivial dependency structure.
Revised: We observe that the performance increases with increasing total latent dimension, likely resulting from the additional ﬂexibility of the higher-dimensional latent space. We also observe that for a ﬁxed number of latent dimensions, models with fewer node dimensions typically perform better. Conversely, models with more nodes, and therefore potentially more complex latent dependency structures, typically have improved performance. This highlights the importance of using an expressive dependency structure for obtaining a ﬂexible model.

Original: 4. Q UANTITATIVE C OMPARISON
Revised: 4. Q UANTITATIVE C OMPARISON

Revised: To quantitatively evaluate the improvements due to learning the latent dependency structure, we compare with a range of common, pre-deﬁned baseline structures.

Revised: These baselines include classic VAEs (Kingma & Welling, 2014; Rezende et al., 2014), which contain no dependencies in the prior, Ladder VAEs (Sønderby et al., 2016), which contain chain-like dependencies in the prior, and fully-connected VAEs (FC-VAEs)

Original: We evaluate the performance of all models using their test-time log-likelihood log p θ ( x ) (Table 1).
Revised: (cf. Kingma et al. (2016)), which contain all possible dependencies in the prior. Quantitative comparisons (Section 4.3) in test log-likelihood demonstrate that the learned dependency structures from the proposed Graph VAE outperform these baselines on each dataset.

Revised: We note that our approach is orthogonal and could be complemented by a number of other approaches attempting to overcome the limitations of classic VAEs.

Revised: Similar to ladder VAEs, Zhao et al.

Revised: (2017) use chain-structured latent dependencies to learn disentangled representations.

Revised: Normalizing ﬂows (Rezende & Mohamed, 2015), on the other hand, adds dependencies to the approximate posterior through a series of invertible transformations.

Revised: We evaluate the performance of all models using their test log-likelihood, log p θ ( x ) (Table 1).

Original: All values were estimated using 5 , 000 importance-weighted samples.
Revised: All values were estimated using 5 , 000 importance-weighted samples.

Original: Following standard practice, we report log p θ ( x ) in nats on MNIST/
Revised: Following standard practice, we report log p θ ( x ) in nats on MNIST/

Original: Omniglot and in bits/input dimension on CIFAR-10. 
Revised: Omniglot and in bits/input dimension on CIFAR-10. The

Original: Our proposed model with a jointly optimized, learned dependency structure  consistently outperformsboth models with less expressive (VAE, ladder VAE) and more expressive (FC-VAE) predeﬁned structures.
Revised: [Table] learned dependency structure in our proposed Graph VAE consistently outperforms models with both fewer (VAE, ladder VAE) and more (FC VAE) latent dependencies.

Original: We discuss potential reasons in section 4.4.
Revised: We discuss potential reasons in Section 5.

Original: To provide further insights into the training objective, Table 1 also reports D KL ( q φ ( z ∣ x )∣∣ 
Revised: To provide further insight into the training objective, Table 1 also reports D KL ( q φ ( z ∣ x )∣∣ p

Original: p θ ( z )) and thefull ELBO objective (Eq. (8)) of the trained models.
Revised: θ ( z )) and the ELBO for each model on the test set.

Original: 4.  D ISCUSSION
Revised:  D ISCUSSION

Original: We want to brieﬂy discuss a number of observations that we have made during training and application of the proposed model and hope they provide the starting point for further research on structure learning in deep generative models.

Original: [Table]

Original: Performance vs. Speed/Memory.
Revised: Performance vs. Speed/Memory.

Original: As shown in Fig.
Revised: As shown in Fig.

Original: 4, the number of latent nodes has a big impacton the performance of our model.
Revised: 4, the number of latent nodes can signiﬁcantly impact the performance of our model.

Original: While allowing more complex dependency structures through a low M / N -ratio is typically beneﬁcial, it also has an adverse effect on the training time and memory consumption.
Revised: While allowing more complex dependency structures through a low M / N -ratio is typically beneﬁcial, it also has an adverse effect on the training time and memory consumption.

Original: Fortunately, the ability to freely select this ratio allows a simple adaption to the available processing power, hardware constraints, and application scenarios.
Revised: Fortunately, the ability to freely select this ratio allows a simple adaption to the available processing power, hardware constraints, and application scenarios.

Original: Optimization Order.
Revised: Optimization Order.

Original: It is worth noting that the learning process optimizes the model parameters ( c ,φ,θ ) in a clear temporal order: While the latent structure, governed by c , converges during the ﬁrst ≈ 200 epochs (Fig. 3a), it takes over 10 × as long until the variational and generative parameters ( φ,θ ) converge.
Revised: It is worth noting that the learning process optimizes the model parameters ( c ,φ,θ ) in a clear temporal order. While the latent structure, governed by c , converges during the ﬁrst ≈ 200 epochs (Fig. 3a), it takes over 10 × as long until the variational and generative parameters ( φ,θ ) converge.

Original: There is no external force incentivizing this behaviour, indicating that the loss can initially most easily be decreased by limiting the latent structure to the complexity prescribed by the observed training data.
Revised: There is no external force enforcing this behaviour, indicating that the loss can initially most easily be decreased by limiting the latent structure to the complexity prescribed by the observed training data.

Original:  Improvement over Fully-Connected VAEs.
Revised: Performance Improvement over Fully-Connected VAEs.

Original: By deﬁnition, the latent spaces of FC-VAEs make no conditional independence assumptions and are therefore more expressive andmore ﬂexible thanthe othermodels in this comparison, including any latent structure resulting from the proposed learning process.
Revised: FC-VAEs model all conditional latent dependencies and are thus potentially more expressive and ﬂexible than other latent dependency structures. It is therefore somewhat surprising that the learned latent structures in Graph VAE consistently outperform the FC-VAE baseline.

Original: The fact that graph VAE outperforms FC-VAE across all datasets is thus interesting, but not entirely surprising: It is, for instance, well-known that latent variable modelsare prone to local optima (Bowman et al., 2016; Burda et al., 2016) and it is reasonable to assume that complex dependencies only add to the optimization challenges.
Revised: We speculate that this may be due to difﬁculties in optimization, which is a known problem in hierarchical latent variable models (Bowman et al., 2016; Burda et al., 2016). Graph VAE and FC-VAE both contain the same hypothesis space of possible models that can be learned. If all dependencies are needed, Graph VAE could set all dependency gate parameters to 1 .

Revised: Likewise, if a latent dependency was unnecessary, FC-VAE could set all of the model parameters in that dependency to 0 .

Revised: However, this would require many coordinated steps along multiple parameter dimensions.

Original: The ability to modify entire latent dependencies allows a data-driven simpliﬁcation of  the underlying structure and, potentially, the optimization task itself.
Revised: It is plausible that the beneﬁt of learning the dependency structure may stem from the ability to alter the optimization landscape, using the dependency gates to move through the model parameter space more coarsely and rapidly. The resulting latent dependency structures may thus be less expressive, but easier to optimize. Follow-up work will be required to test this intuition.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: We presented a novel method for optimizing variational autoencoders jointly with their latent  dependency structures. Our experiments showed that the learned latent dependency structure improves thegenerative performance of latent variable models. By introducing an additional set of binary structure indicator variables and optimizing a structure-dependent evidence lower bound, our model is able to learn a better representation of the underlying generative distribution than baselines with a predetermined structure.
Revised: We presented a novel method for structure learning in latent variable models, which uses dependency gating variables, together with a modiﬁed objective and discrete differentiation techniques, to effectively transform discrete structure learning into a continuous optimization problem. In our experiments, the learned latent dependency structures improve the performance of latent variable models over comparable baselines with predeﬁned dependency structures.

Original: We believe that the proposed model is useful beyond the applications shown in this work and can be extended inmany interesting ways. Examples include structure learning on undirected Markov random ﬁelds, dynamic structure prediction in time-series models, and joint optimization with the neural encoder-decoder architecture.
Revised: The approach presented here provides directions for further research in structure learning for other tasks, including undirected graphical models, time-series models, and discriminative models.

Original: AS UPPLEMENTARY M ATERIAL A. L OWER B OUND D ERIVATION
Revised: A L OWER B OUND D ERIVATION

Original: Introducing the  dependency gating variables,  c  , modiﬁes the evidence lower bound (ELBO) from Eq.
Revised: Introducing the distribution over dependency gating variables, p ( c ) = B( µ ) , modiﬁes the evidence lower bound (ELBO) from Eq.

Original: (2), as we now have an additional set of random variables over which to marginalize.
Revised: (2), as we now have an additional set of random variables over which to marginalize.

Original: To see this, we can start by re-expressing Eq.
Revised: To see this, we can start by re-expressing Eq.

Original: (2) as
Revised: (2) as

Original: [Equation]
Revised: [Equation]

Original: p θ ( x , z ) can be expressed as a marginalization over the gating variables, effectively averaging over an ensemble of models with a distribution of dependency structures:
Revised: p θ ( x , z ) can be expressed as a marginalization over the gating variables, effectively averaging over an ensemble of models with a distribution of dependency structures:

Original: [Equation]
Revised: [Equation]

Original: Plugging this into Eq.
Revised: Plugging this into Eq.

Original: [Equation]
Revised: [Equation]

Original: Using Jensen’s inequality
Revised: Using Jensen’s inequality

Original: [Equation]
Revised: [Equation]

Original: wecan bring the log inside of the expectationE p ( c ) [⋅]
Revised: we bring the log inside of the expectation

Original: [Equation]
Revised: [Equation]

Original:  yielding
Revised: E p ( c ) [⋅] , yielding

Original: [Equation]
Revised: [Equation]

Original: ) wherewe have deﬁned ̃L as a lower bound on L (and therefore still a lower bound on log p θ ( x ) ).
Revised: where ̃L is a lower bound on L .

Original: Swapping the order of expectation, wecan rewrite ̃L as
Revised: Swapping the order of expectation [Equation] we rewrite ̃L as

Original: [Equation]
Revised: [Equation]

Original: and because the second term is independent of p ( c ) , we can include both terms inside of a single outer expectation:
Revised: and because the second term is independent of p ( c ) , we can include both terms inside of a single outer expectation:

Original: [Equation]
Revised: [Equation]

Original: where we have deﬁned L c as the ELBO for a given dependency structure.
Revised: where we have deﬁned L c as the ELBO for a given dependency structure.

Revised: Note that ̃L is not a proper variational bound, as it cannot recover the marginal log likelihood.

Revised: Rather, ̃L allows us to optimize the distribution over gating variables and, thus, the model structure.

Revised: When we arrive at a ﬁxed structure, we will be optimizing the variational bound for that particular dependency structure.

Original: Thus, during training, we can estimate ̃L by sampling gating variables, c ∼ p ( c )  , and evaluating Monte Carlo estimates of
Revised: To see this, note that the bound in Eq. 13 becomes tight when p ( c ) is any ﬁxed distribution, in which the dependency gating means, µ , are all either 1 or 0 . Plugging in a delta distribution for p ( c ) at a particular conﬁguration, c ′ , i.e. p ( c ) = δ c , c ′ , we have

Revised: [Equation]

Revised: Intuitively, this is simply the case of a latent variable model with predeﬁned, ﬁxed dependencies.

Original: L c  .
Revised: By optimizing ̃L w.r.t. µ , we hope to collapse p ( c ) to a delta distribution at the optimal (MAP) conﬁguration, c ∗ , because [Equation] Effectively, we can search for the dependency structure with the highest ELBO by optimizing the distribution parameters of p ( c ) .

Revised: Although this optimization procedure may arrive at locally optimal dependency structures, our hope is that these learned structures will still perform better than an arbitrary, predeﬁned dependency structure.

Original: A. I MPLEMENTATION D ETAILS A.2. N ETWORK A RCHITECTURE
Revised: B I MPLEMENTATION D ETAILS B. N ETWORK A RCHITECTURE

Original: We document the network architectures used in our experiments.
Revised: We document the network architectures used in our experiments.

Original: We use the same network architectures for all datasets.
Revised: We use the same network architectures for all datasets.

Original: Input size (pixels per image) is the only difference across datasets.
Revised: Input size (pixels per image) is the only difference across datasets.

Original: The input size is 28 × 28 for MNIST and Omniglot, and 3 × 32 × 32 for CIFAR-10.
Revised: The input size is 28 × 28 for MNIST and Omniglot, and 3 × 32 × 32 for CIFAR-10.

Original: Encoders:
Revised: Encoders:

Revised: fc(input size

Original: [Equation]
Revised: [Equation]

Revised: 512) → Batch Norm → ELU

Revised: [Equation]

Revised: [Equation]

Revised: Batch Norm

Revised: [Equation]

Revised: ELU

Revised: [Equation]

Revised: 512, 512) → Batch Norm → ELU → fc(512, 256)

Original: Latent Dependencies: Latent dependencies are modelled by non-linear MLPs.
Revised: Latent Dependencies: Latent dependencies are modelled by non-linear MLPs.

Original: Note that the topdown architecture is shared between inference model and generative model.
Revised: Note that the topdown architecture is shared between inference model and generative model.

Original: But the MLPs are optimized independently.
Revised: But the MLPs are optimized independently.

Original: bottom-up: For each node with N ′ dimensions, the local potential is predicted by a mapping from the encoded feature:
Revised: bottom-up: For each node with N ′ dimensions, the local potential is predicted by a mapping from the encoded feature:

Original: fc( ) → Batch Norm → ReLU → fc(256, N ′ ).
Revised: fc(256, 256) → Batch Norm → ReLU → fc(256, N ′ ).

Original: The output feature is then mapped to µ and log var with two independent fc( N ′ , N ′ ) respectively.
Revised: The output feature is then mapped to µ and log var with two independent fc( N ′ , N ′ ) respectively.

Original: top-down: For each node ( N ′ dimensions) with a set of parent nodes, the top-down inference/generation is implemented as:
Revised: top-down: For each node ( N ′ dimensions) with a set of parent nodes, the top-down inference/gen-

Original: fc(sum
Revised: eration is

Original:  of parent nodes’ dimension → Batch Norm → ReLU → fc(256, N ′ )
Revised: implemented as: fc(sum of parent nodes’ dimension, 256) → Batch Norm → ReLU → fc(256, N ′ )

Original: The output feature is then mapped to µ and log var with two independent fc( N ′ , N ′ ) respectively.
Revised: The output feature is then mapped to µ and log var with two independent fc( N ′ , N ′ ) respectively.

Original: Decoders:
Revised: Decoders:

Original: fc

Original: [Equation]

Original: [Equation]

Original: Batch Norm

Original: [Equation]

Original: ReLU

Original: [Equation]

Original: 512, 512) → Batch Norm → ReLU → fc(512, 512) → Batch Norm → ReLU → fc(512, input size) → out put function() out put function for MNIST and Omniglot is sigmoid() which predicts the µ of Bernoulli observations ; and sigmoid predicting µ , fc(input size, input size) predicting log var of Gaussion observations for CIFAR-10.
Revised: fc(512, 512) → Batch Norm → ReLU → fc(512, 512) → Batch Norm → ReLU → fc(512, 512) → Batch Norm → ReLU → fc(512, input size) → out put function() out put function for MNIST and Omniglot is sigmoid() which predicts the µ of Bernoulli observations; and sigmoid predicting µ , fc(input size, input size) predicting log var of Gaussion observations for CIFAR-10.

Original: A.2. T RAININGAll
Revised: B. T RAINING

Original:  models were implemented with PyTorch (Paszke et al., 2017) and trained using the Adam (Kingma & Ba, 2015) optimizer with a mini-batch size of 128 , and learning rate of 3 e − 4 .
Revised: All models were implemented with PyTorch (Paszke et al., 2017) and trained using the Adam (Kingma & Ba, 2015) optimizer with a mini-batch size of 128 , and learning rate of 3 e − 4 .

Original: The
Revised: The

Original: Gumbel-softmax temperature were initialized as 1 , and decrease to 0 .
Revised: Gumbel-softmax temperature were initialized as 1 , and decrease to 0 .

Original: 999 epoch at each epoch.
Revised: 99 epoch at each epoch.

Original: MNIST and Omniglot took 2000 epochs to converge, and CIFAR took 3000 epochs to converge. 
Revised: MNIST and Omniglot took 2000 epochs to converge, and CIFAR took 3000 epochs to converge. B. I NFERENCE M ODULE D ETAILS

Revised: Parameter prediction for the local factors q φ ( z n ∣ x , z pa ( n ) ) consists of a precision-weighted fusion of the top-down prediction ̂ ψ n and a bottom-up prediction ̃ ψ n .

Revised: Speciﬁcally, for a latent variable z n , ̂ ψ n ∶= {̂ µ n , ̂ σ n } , and ̃ ψ n ∶= {̃ µ n , ̃ σ n } .

Revised: Bottom-up Inference.

Revised: A high-dimensional input is ﬁrst mapped to a feature vector h x by an encoder

Revised: MLP.

Revised: h x is then used to predict ̃ µ n and ̃ σ n with non-linear MLP BUn .

Revised: Top-down Inference.

Revised: ̂ µ n and ̂ σ n are predicted by ̂ µ n = MLP TDn ([ z pa ( n ) ⊙ c pa ( n ) ,n ]) and

Revised: ̂ σ n =

Revised: MLP TDn ([ z pa ( n ) ⊙ c pa ( n ) ,n ]) , respectively.

Revised: [,] denotes concatenation operation, and ⊙ denotes element-wise multiplication.

Revised: Precision-weighted fusion.

Revised: Having ̂ ψ n and ̃ ψ n , the parameters of the local conditional distribution is given by q φ ( z n ∣ x , z pa ( n ) , c pa 9 n ) ,n ) ∼ N( z n ∣ µ n ,σ 2 n ) , with

Revised: [Equation]

Revised: C A DDITIONAL R ESULTS C. L OG - LIKELIHOOD C OMPARISON

Revised: Test log-likelihood comparison at different epochs is given in Fig.

Revised: Averaged log-likelihoods and associated standard deviations are calculated from 5 individual runs.

Revised: [Figure]

Revised: Although we do not claim that Graph VAE learns a disentangled/semantically meaningful latent space, we do agree this is an interesting and potentially useful direction.

Revised: In this section, we plot the TSNE visualization of the data as well as latent embeddings of Graph VAE, VAE on MNIST in Fig.

Revised: As shown, The latent embedding of Graph VAE exhibits a large gap between different classes, even though the model is trained in an unsupervised fashion.

Revised: We will further investigate on this perspective in future work.

Revised: [Figure]