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 - work combines the complementary strengths of deep generative models and probabilistic graphical models.
Revised: Our general modeling and inference framework 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 singlevariational objective.
Revised: The network parameters, variational parameters as well as the latent topology are optimized simultaneously with a single 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 an 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 result in models that more accurately model the data distribution, outperforming several common predeﬁned latent dependency structures.
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]

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; Goyal et al., 2017; Vikram et al., 2018; Webb et al., 2018).
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 θ [Equation] z m ∣ z pa ( m ) )
Revised: While these conditional distributions may be relatively simple, the marginal empirical prior, p θ ( z m ) = ∫ p θ ( 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: 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: The discrete nature of the latent space’s topological structure introduces discontinuities in the optimization landscape, complicating the learning process.
Revised: The discrete nature of the latent space’s topological structure introduces discontinuities in the optimization landscape, complicating the learning process.

Original: 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:
Revised: 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 an appropriate 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.

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

Original: 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).
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).

Original: 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).
Revised: 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).

Original: Indeed, we observe this process empirically, with the learned dependency structures outperforming their predeﬁned counterparts (Section 4).
Revised: 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: 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: Each possible latent dependency structure can now be expressed through its corresponding value of c .
Revised: Each possible latent dependency structure can now be expressed through its corresponding value of c .

Original: Each ﬁxed latent dependency structure, c ′ , induces a separate latent variable model p θ ( x , z , c )
Revised: Each ﬁxed latent dependency structure, c ′ , induces a separate latent variable model p θ ( x , z , c )

Original: = p θ ( x ∣ z , c ) p θ ( z ∣ c ) δ c , c ′ , where δ ⋅ , ⋅ is the Kronecker delta, which effectively selects a single structure.
Revised: = 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 on the latent variables 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) 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
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: B for additional details on the MLP architecture).
Revised: B for additional details on the MLP architecture).

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

Original: N( 0 , I ) .
Revised: 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 , c ) , must approximate p θ ( z ∣ x , c ) .
Revised: The approximate posterior, q φ ( z ∣ x , c ) , must approximate p θ ( z ∣ x , c ) .

Original: We express the approximate posterior as
Revised: We express the approximate posterior as

Original: [Equation]
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 .
Revised: -conditioned parameters ̃ ψ n from a generic encoding of x .

Original: While Webb et al. (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.
Revised: We note that the dependency structures of the generative model and its corresponding posterior are, in general, not independent: choosing a particular structure in the generative model induces a particular structure in the posterior (Webb et al., 2018).

Revised: A simple way to guarantee enough capacity in the encoder to account for the dependencies implied by the decoder is thus to keep the encoder graph fully-connected and learn a decoder graph only.

Original: This reduces the number of learned parameters, and  we found that it gave improved empirical performance  over maintaining all dependencies.
Revised: Instead, we share the gating variables c between approximate posterior and generative model (see Section 3.3), i.e., we assume that the encoder dependencies mirror those of the decoder. As a consequence, the posterior implied by the generative model could lie outside of the model class representable by the encoder. In practice, this is not an issue and we observe signiﬁcant performance improvements over both traditional VAEs (where prior and posterior match but are limited in their expressiveness) and Graph VAEs with fully-connected encoder graph. See Section 4.3 for quantitative experiments and Section 5 for a discussion on the relationship between learned structures and fully-connected structures.

Original: 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 ̂
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: Additional details on the fusion process can be found in Appendix B.3.
Revised: Additional details on the fusion process can be found in Appendix B.3.

Original: 3. L EARNING S TRUCTURE BY I NDUCING AN A DDITIONAL L OWER B OUND
Revised: 3. L EARNING S TRUCTURE BY I NDUCING AN A DDITIONAL L OWER B OUND

Original: The formulation in Section 3.1 provides the form of the model for a particular conﬁguration of the latent dependency structure.
Revised: The formulation in Section 3.1 provides the form of the model for a particular conﬁguration of the latent dependency structure.

Original: Finding the optimal structure corresponds to a discrete optimization over all values of c , potentially optimizing the model parameters of each possible conﬁguration.
Revised: 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: 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.
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.

Original: Speciﬁcally, we treat each c i,j as an independent random variable, sampled from a Bernoulli distribution with mean µ i,j ,
Revised: Speciﬁcally, 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 ∼ p ( c i,j ) =
Revised: i.e. c i,j ∼ p ( c i,j ) =

Original: B( µ i,j ) .
Revised: B( µ i,j ) .

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

Original: Introducing this distribution allows us to express the following additional lower bound on L , derived in Appendix A:
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 particular value of dependency gating variables.
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 ) , which we estimate by sampling c ∼ p ( c ) and evaluating L 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 .

Original: 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.
Revised: We note that ̃L is not a proper variational bound, as it is not guaranteed to recover the marginal likelihood if the approximate posterior Require: Data x , number of latent nodes N , dimensions per node N matches the true posterior.

Original: Rather, optimizing ̃L provides a method for learning the latent structure.
Revised: Rather, optimizing ̃L provides a method for learning the latent structure.

Original: For ﬁxed parameters ( φ,θ ) , the optimal ̃ L w.r.t.
Revised: For ﬁxed parameters ( φ,θ ) , the optimal ̃L w.r.t.

Original: µ is a δ -distribution at the MAP conﬁguration, c ∗ , yielding L = ̃L
Revised: µ is a δ -distribution at the MAP conﬁguration, c ∗ , yielding L = ̃L

Original: =
Revised: =

Original: L c ∗ .
Revised: L c ∗ .

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

Original: In practice, we jointly optimize φ , θ , and µ .
Revised: In practice, we jointly optimize φ , θ , and µ .

Original: 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).
Revised: 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).

Original: Thus, by the end of training, we are effectively optimizing the ELBO for a single dependency structure.
Revised: Thus, by the end of training, we are effectively optimizing the ELBO for a single dependency structure.

Original: The training procedure is outlined in Algorithm 1.
Revised: 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: (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.
Revised: 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 proposed latent dependency learning approach on three benchmark datasets:
Revised: We evaluate the proposed latent dependency learning approach on three benchmark datasets:

Original: MNIST (Lecun et al., 1998; Larochelle & Murray, 2011), Omniglot (Lake et al., 2013), and CIFAR10 (Krizhevsky, 2009).
Revised: MNIST (Lecun et al., 1998; Larochelle & Murray, 2011), Omniglot (Lake et al., 2013), and CIFAR10 (Krizhevsky, 2009).

Original: 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.
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: In Section 4.3, we provide quantitative comparisons with common predeﬁned latent dependency structures on benchmark datasets.
Revised: In Section 4.3, we provide quantitative comparisons with common predeﬁned latent dependency structures on benchmark datasets.

Original: Additional results on training robustness and latent space embeddings can be found in Appendix C.
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 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.
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 models and experiments ( M = 80 ).
Revised: Likewise, the number of [Figure] latent dimensions is the same in all models and experiments ( M = 80 ).

Original: As discussed in Section 3.1, all latent dependencies are modeled 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 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.
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 p θ ( x ∣ z ) with a Gaussian density, with mean and log-variance predicted by sigmoid and linear functions, respectively.
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.

Original: 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. [Figure]
Revised: Further implementation details, including the model architectures and training criteria, can be found in Appendix B.

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

Original: We ﬁrst explore the structure learning process.
Revised: We ﬁrst explore the structure learning process.

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

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

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

Original: In Fig.
Revised: In Fig.

Original: 3, we visualize this process during training on MNIST.
Revised: 3, we visualize this process during training on MNIST.

Original: The model has N = 5 nodes and a total latent dimension of M = [-
 (i e.,-] N ′ =
Revised: The model has N = 5 nodes and a total latent dimension of M = [-
 (i e.,-]80 (i.e., N ′ =

Original: M / N = 16 dimensions per node).
Revised: M / N = 16 dimensions per node).

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

Original: 3a, the gating means converge in practice with 3 out of 10 edges removed and the rest retained.
Revised: 3a, the gating means converge in practice, with 3 out of 10 edges removed and the rest retained.

Original: The resulting static dependency structure is visualized in Fig.
Revised: The resulting static dependency structure is visualized in Fig.

Original: 3b.
Revised: 3b.

Original: 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.
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: 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 ′ =
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 .

Original: Results are shown in Fig.
Revised: Results are shown in Fig.

Original: 4 for various models trained on MNIST.
Revised: 4 for various models trained on MNIST.

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

Original: We observe that the performance increases with increasing total latent dimension, likely resulting from the additional ﬂexibility of the higher-dimensional latent space.
Revised: We observe that the performance increases with increasing total latent dimension, likely resulting from the additional ﬂexibility of the higher-dimensional latent space.

Original: We also observe that for a ﬁxed number of latent dimensions, models with fewer node dimensions typically perform better.
Revised: We also observe that for a ﬁxed number of latent dimensions, models with fewer node dimensions typically perform better.

Original: Conversely, models with more nodes, and therefore potentially more complex latent dependency structures, typically have improved performance.
Revised: Conversely, models with more nodes, and therefore potentially more complex latent dependency structures, typically have improved performance.

Original: This highlights the importance of using an expressive dependency structure for obtaining a ﬂexible model.
Revised: 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

Original: 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: To quantitatively evaluate the improvements due to learning the latent dependency structure, we compare with a range of common, predeﬁned baseline structures.

Original: 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) (cf.
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 fullyconnected VAEs (FC-VAEs) (cf.

Original: Kingma et al.
Revised: Kingma et al.

Original: (2016)), which contain all possible dependencies inthe prior.
Revised: (2016)), which contain all possible dependencies in

Original: 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: [Table] the prior (corresponding to all gating variable parameters µ i,j set to a ﬁxed value of 1 ).

Original: 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: 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.

Original: Similar to ladder VAEs, Zhao et al.
Revised: Similar to ladder VAEs, Zhao et al.

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

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

Original: We evaluate the performance of all models using their test log-likelihood, log p θ ( x ) (Table 1).
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.

Original: The[Table] learned dependency structure in our proposed Graph VAE consistently outperforms models with both fewer (VAE, ladder VAE) and more (FC VAE) latent dependencies.
Revised: The 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 5.
Revised: We discuss potential reasons in Section 5.

Original: To provide further insight into the training objective, Table 1 also reports
Revised: To provide further insight into the training objective, Table 1 also reports

Original: D KL ( q φ ( z ∣ x )∣∣ p
Revised: D KL ( q φ ( z ∣ x )∣∣ p

Original: θ ( z )) and the ELBO for each model on the test set.
Revised: θ ( z )) and the ELBO for each model on the test set.

Original: D ISCUSSION
Revised: Encoder-Decoder Relationship. From a purely theoretical point of view, learning the structure of the generative model implies the need for a fully-connected graph in the approximate posterior (see Section 3.1). In practice, we share the gating variables c between encoder and decoder (see Eq.

Revised: (8)), because we observed improved empirical performance when doing so: training a Graph VAE decoder with a predeﬁned, fully-connected encoder graph results in a test log-likelihood of − 83 .

Revised: 4 nats on MNIST, which is closer to the performance of FC-VAE ( − 83 . 0 nats) than Graph VAE ( − 82 .nats).

Revised: The KL-term and ELBO in this experiment are 26 .

Revised: 5 and − 85 .

Revised: 1 , respectively.

Revised: We believe these are noteworthy empirical results, but further research is required to understand this behavior at a theoretical level.

Revised: A full visualization of the training process is provided in Appendix C.1.

Revised:  D ISCUSSION

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 can signiﬁcantly impact 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.
Revised: It is worth noting that the learning process optimizes the model parameters ( c ,φ,θ ) in a clear temporal order.

Original: 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: 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 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.
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: Performance Improvement over Fully-Connected VAEs.
Revised: Performance Improvement over Fully-Connected VAEs.

Original:  FC-VAEs model all conditional latent dependencies and are thus potentially more expressive and ﬂexible than other latent dependency structures.
Revised: There is an intricate relationship between fully-connected graphs vs. learned graphs along one axis and between prior structures vs. posterior structures along a separate axis: FC-VAEs model all conditional latent dependencies and are thus potentially more expressive and ﬂexible than other latent dependency structures.

Original: It is therefore somewhat surprising that the learned latent structures in Graph VAE consistently outperform the FC-VAE baseline.
Revised: It is therefore somewhat surprising that the learned latent structures in Graph VAE consistently outperform the FC-VAE baseline.

Original: 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).
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).

Original: Graph VAE and FC-VAE both contain the same hypothesis space of possible models that can be learned.
Revised: Graph VAE and FC-VAE both contain the same hypothesis space of possible models that can be learned.

Original: If all dependencies are needed, Graph VAE could set all dependency gate parameters to 1 .
Revised: If all dependencies are needed, Graph VAE could set all dependency gate parameters to 1 .

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

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

Original: 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.
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.

Original: The resulting latent dependency structures may thus be less expressive, but easier to optimize.
Revised: The resulting latent dependency structures may thus be less expressive, but easier to optimize.

Revised: While only an intuition, this hypothesis is also in line with the observations and results in our experiments with fully-connected encoder graphs and learned decoder graphs, where the theoretically more ﬂexible FC-encoder is outperformed by our parameter sharing approach.

Original: Follow-up work will be required to test this intuition.
Revised: Follow-up work will be required to test this intuition.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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.
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.

Original: In our experiments, the learned latent dependency structures improve the performance of latent variable models over comparable baselines with predeﬁned dependency structures.
Revised: In our experiments, the learned latent dependency structures improve the performance of latent variable models over comparable baselines with predeﬁned dependency structures.

Original: 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
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: A L OWER B OUND D ERIVATION
Revised: A L OWER B OUND D ERIVATION

Original: Introducing the distribution over dependency gating variables, p ( c )
Revised: Introducing the distribution over dependency gating variables, p ( c )

Original: =
Revised: =

Original: B( µ ) , modiﬁes the evidence lower bound (ELBO) from Eq.
Revised: 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: we bring the log inside of the expectation
Revised: we bring the log inside of the expectation

Original: [Equation]
Revised: [Equation]

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

Original: [Equation]
Revised: [Equation]

Original: where ̃L is a lower bound on L .
Revised: where ̃L is a lower bound on L .

Original: Swapping the order of expectation
Revised: Swapping the order of expectation

Original: [Equation]
Revised: [Equation]

Original: we rewrite ̃L as
Revised: 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.

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

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

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

Original: To see this, note that the bound in Eq.
Revised: To see this, note that the bound in Eq.

Original: 13 becomes tight when p ( c ) is any ﬁxed distribution, in which the dependency gating means, µ , are all either 1 or 0 .
Revised: 13 becomes tight when p ( c ) is any ﬁxed distribution, in which the dependency gating means, µ , are all either 1 or 0 .

Original: Plugging in a delta distribution for p ( c ) at a particular conﬁguration, c ′ , i.e. p ( c ) =
Revised: Plugging in a delta distribution for p ( c ) at a particular conﬁguration, c ′ , i.e. p ( c ) =

Original: δ c , c ′ , we have
Revised: δ c , c ′ , we have

Original: [Equation]
Revised: [Equation]

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

Original: By optimizing ̃L w.r.t.
Revised: By optimizing ̃L w.r.t.

Original: µ , we hope to collapse p ( c ) to a delta distribution at the optimal (MAP) conﬁguration, c ∗ , because
Revised: µ , we hope to collapse p ( c ) to a delta distribution at the optimal (MAP) conﬁguration, c ∗ , because

Original: [Equation]
Revised: [Equation]

Original: Effectively, we can search for the dependency structure with the highest ELBO by optimizing the distribution parameters of p ( c ) .
Revised: Effectively, we can search for the dependency structure with the highest ELBO by optimizing the distribution parameters of p ( c ) .

Original: 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.
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: B I MPLEMENTATION D ETAILS B. 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:

Original: fc(input size
Revised: fc(input size

Original: [Equation]
Revised: [Equation]

Original: 512) → Batch Norm → ELU
Revised: 512) → Batch Norm → ELU

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Batch Norm
Revised: Batch Norm

Original: [Equation]
Revised: [Equation]

Original: ELU
Revised: ELU

Original: [Equation]
Revised: [Equation]

Original: 512, 512) → Batch Norm → ELU → fc(512, 256)
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:
Revised: bottom-up:

Original: For each
Revised: For each

Original: node with
Revised: node with

Original: N ′ dimensions, the local potential is predicted by a mapping from the encoded feature:
Revised: N ′ dimensions, the local potential is predicted by a mapping from the encoded feature:

Original: fc(256, 256) → 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/gen-
Revised: top-down: For each node ( N ′ dimensions) with a set of parent nodes, the top-down inference/gen-

Original: eration is
Revised: eration is

Original: implemented as: fc(sum of parent nodes’ dimension, 256) → 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(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.
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: B. T RAINING
Revised: B. T RAINING

Original: 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 .
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: 99 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.

Original: B. I NFERENCE M ODULE D ETAILS
Revised: B. I NFERENCE M ODULE D ETAILS

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 ̂ ψ n and a bottom-up prediction ̃ ψ n .
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 .

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

Original: Bottom-up Inference.
Revised: Bottom-up Inference.

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

Original: MLP.
Revised: MLP.

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

Original: Top-down Inference.
Revised: Top-down Inference.

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

Original: ̂ σ n =
Revised: ̂ σ n =

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

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

Original: Precision-weighted fusion.
Revised: Precision-weighted fusion.

Original: 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: Having ̂ ψ n and ̃ ψ n , the parameters of the local conditional distribution is given by q φ ( z n ∣ x , z pa ( n ) , c pa ( n ) ,n ) ∼ N( z n ∣ µ n ,σ 2 n ) , with

Original: [Equation]
Revised: [Equation]

Original: C A DDITIONAL R ESULTS C.  L OG - LIKELIHOOD C OMPARISON
Revised: C A DDITIONAL R ESULTS C. R OBUSTNESS OF L OG - LIKELIHOODS

Revised: In Fig.

Original: Test log-likelihood comparison at different epochs is given in Fig. Averaged log-likelihoods and associated standard deviations  are calculated from 5 individual runs.
Revised: 5, we report the averaged test log-likelihoods and associated standard deviations of Graph VAE and our baselines at different epochs. All calculations are based on 5 independent runs.

Original: [Figure] 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: Our training objective optimizes intrinsic structure (which does not necessarily correlate with semantic meaning) and does not incentivize a disentanglement of latent factors.

Original: In this section, we plot the TSNE visualization of the data as well as latent embeddings of Graph VAE, VAE on MNIST in Fig. 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: Interestingly, a TSNEvisualization of the data as well as latent embeddings of Graph VAE and VAE on MNIST (Fig. 6) shows that the latent embedding of Graph VAE exhibits a large (semantic) gap between different classes, even though the model is trained in an unsupervised fashion.

Original: We will further investigateon this perspective in future work.
Revised: We will further investigate this behavior in future work.

Original: [Figure]
Revised: [Figure]