Title: Structured Neural Networks for Density Estimation and Causal Inference
Venue: NeurIPS 2023
Accepted: yes
Keywords: generative models

density estimation

normalizing flows

binary matrix factorization

causal inference
Abstract: Injecting structure into neural networks enables learning functions that satisfy invariances with respect to subsets of inputs. For instance, when learning generative models using neural networks, it is advantageous to encode the conditional independence structure of observed variables, often in the form of Bayesian networks. We propose the Structured Neural Network (StrNN), which injects structure through masking pathways in a neural network. The masks are designed via a novel relationship we explore between neural network architectures and binary matrix factorization, to ensure that the desired independencies are respected. We devise and study practical algorithms for this otherwise NP-hard design problem based on novel objectives that control the model architecture. We demonstrate the utility of StrNN in three applications: (1) binary and Gaussian density estimation with StrNN, (2) real-valued density estimation with Structured Autoregressive Flows (StrAFs) and Structured Continuous Normalizing Flows (StrCNF), and (3) interventional and counterfactual analysis with StrAFs for causal inference. Our work opens up new avenues for learning neural networks that enable data-efficient generative modeling and the use of normalizing flows for causal effect estimation.

Structured Neural Networks for Density Estimation and Causal Inference

Abstract

Injecting structure into neural networks enables learning functions that satisfy invariances with respect to subsets of inputs. For instance, when learning generative models using neural networks, it is advantageous to encode the conditional independence structure of observed variables, often in the form of Bayesian networks. We propose the Structured Neural Network (StrNN), which injects structure through masking pathways in a neural network. The masks are designed via a novel relationship we explore between neural network architectures and binary matrix factorization, to ensure that the desired independencies are respected. We devise and study practical algorithms for this otherwise NP-hard design problem based on novel objectives that control the model architecture. We demonstrate the utility of StrNN in three applications: (1) binary and Gaussian density estimation with StrNN, (2) real-valued density estimation with Structured Autoregressive Flows (StrAFs), autoregressive normalizing flows that leverage StrNN as a conditioner, and (3) interventional and counterfactual analysis with StrAFs. Our work opens up new avenues for learning neural networks that enable data-efficient generative modeling and the use of normalizing flows for causal effect estimation.

1 Introduction

The incorporation of structure into machine learning models has been shown to provide benefits for model generalization, learning efficiency, and interpretability. The improvements are particularly salient when learning from small amounts of data. This idea has found use in reinforcement learning [Ok et al., 2018], computational healthcare [Hussain et al., 2021, Cui et al., 2020], survival analysis [Gharari et al., 2023], time series analysis [Curi et al., 2020], and causal inference [Balazadeh et al., 2022].

This work focuses on the problem of density estimation from high-dimensional data which has been approached through a variety of lenses. For example, normalizing flows [Tabak and Turner, 2013, Rezende and Mohamed, 2015] model data by transforming a base distribution through a series of invertible transformations. Masked autoencoders (MADE) [Germain et al., 2015] model the joint distribution via an autoregressive factorization of

Weight

Mask

Weight

Mask

Adjacency

Matrix

Weight

Mask

the random variables. The factorization is enforced by integrating structure in the neural network of the autoencoder. The MADE architecture zeros out weights in a neural network to ensure each output dimension has an autoregressive dependence on the input dimensions. When data is scarce, this may lead to over-fitting and harm generalization. When knowledge of a Bayesian network [Pearl, 2011] and the associated conditional independencies exist, it is desirable to inject this knowledge directly into the network to improve density estimation. In this work, we extend the concept of weight masking, as in MADEs, to go beyond autoregressive dependencies of the output on the input.

This work proposes the Structured Neural Network (StrNN) , a network architecture that enforces functional independence relationships between inputs and outputs via weight masking. In other words, the output can remain unaffected by changes to (subsets of) the input. We focus on instantiating this idea to model conditional independence between inputs when neural networks are deployed to model the density of random variables, as illustrated in Figure 1. Any set of conditional independence statements (e.g. in a Bayesian network) may be represented via a binary adjacency matrix. StrNN performs binary matrix factorization to generate a set of weight masks that follow the adjacency matrix. There are two key challenges we overcome. First, the general problem of binary matrix factorization in this context is under-specified, as there exist many valid masks whose matrix product realizes a given adjacency matrix. To this end, we propose the idea of neural network path maximization as a strategy to guide the generation of optimal masks. Secondly, binary matrix factorization is NP-hard in general. We study practical solutions that generate valid mask matrices efficiently.

StrNN can then be applied to NN-based density estimation in various contexts. Where conditional independence properties are known a priori, we show that StrNN can be used to estimate parameters of data distributions while keeping specified variables conditionally independent. We further integrate StrNNs into various discrete and continuous flow architectures, including the autoregressive normalizing flow [Papamakarios et al., 2017, Huang et al., 2018] to form the Structured Autoregressive Flow (StrAF) . The StrAF model uses the StrNN as a normalizing flow conditioner network, thus enforcing a given adjacency structure within each flow layer. The StrAF preserves variable orderings between chained layers, allowing the adjacency structure to be respected throughout the entire flow.

Finally, we study a natural application of StrAF in causal effect estimation, a domain that often requires flexible methods for density estimation that enforce conditional independence relationships within causal graphs [Pearl, 2009]. We show how StrAFs can be used to perform interventional and counterfactual queries better than existing flow-based causal models that do not incorporate graphical structures. Across the board, we highlight how incorporating conditional independence structure improves generalization error when learning from a small number of samples.

In summary, the main contributions of this work are as follows:

We introduce StrNN, a weight-masked neural network that can efficiently learn functions with specific variable dependence structures. In particular, it can inject prior domain knowledge in the form of Bayesian networks to probability distributions. We formalize the weight masking as an optimization problem, where we can pick the objective based on desired neural architectures. We propose an efficient binary matrix factorization algorithm to mask arbitrary neural networks.

We integrate StrNN into autoregressive and continuous normalizing flows for best-in-class performance in density estimation and sample generation tasks.

We apply StrAF for causal effect estimation and showcase its ability to outperform existing causal flow models in accurately addressing interventional and counterfactual queries.

2 Background

Masked Autoencoders for Density Estimation (MADE): Masked neural networks were introduced for density estimation on binary-valued data [Germain et al., 2015]. Given x = ( x 1 , ..., x d ) , MADE factorizes p ( x ) as the product of the outputs of a neural network. Writing the j -th output as the conditional probability ˆ x j := p ( x j = 1 | x <j ) , the joint distribution can be rewritten exactly as the binary cross-entropy loss. As long as the neural network outputs are autoregressive in relation to its inputs, we can minimize the cross-entropy loss for density estimation. To enforce the autoregressive property for a neural network y = f ( x ) with a single hidden layer and d inputs and outputs, MADE element-wise multiplies the weight matrices W and V with binary masks M W and M V :

The autoregressive property is satisfied as long as the product of the masks, M V M W ∈ R d × d , is lower triangular. The MADE masking algorithm (2) can be extended to neural networks with an arbitrary number of hidden layers and hidden sizes. For Gaussian data, the MADE model can be extended as R d → R 2 d , ( x 1 , ...x d ) → (ˆ µ 1 , ..., ˆ µ d , log(ˆ σ 1 ) , ..., log(ˆ σ d )) . The last mask must be duplicated to ensure µ j and σ j only depend on x <j . MADE can also be used as the conditioner in an autoregressive flow to model general data, as seen in Papamakarios et al. [2017].

Normalizing Flows: Normalizing flows [Rezende and Mohamed, 2015] model complex data distributions and have been applied in many scenarios [Papamakarios et al., 2021]. Given data x ∈ R d , a normalizing flow T : R d → R d takes x to latent variables z ∈ R d that are distributed according to a simple base distribution p z , such as the standard normal. The transformation T must be a diffeomorphism (i.e., differentiable and invertible) so that we can compute the density of x via the change-of-variables formula: p x ( x ) = p z ( T ( x )) | det J T ( x ) | . We can compose multiple diffeomorphic transformations T k to form the flow T = T 1 ◦· · · ◦ T K since diffeomorphisms are closed under composition. The flows are trained by maximizing the log-likelihood of the observed data under the density p x ( x ) . The log-likelihood can be evaluated efficiently when it is tractable to compute the Jacobian determinant of T ; for example when T k is a lower triangular function [Marzouk et al., 2016]. Given the map, we can easily generate i.i.d. samples from the learned distribution by sampling from the base distribution z i ∼ p z and evaluating the flow T -1 ( z i ) .

Density Estimation with Autoregressive Flows: When the Jacobian matrix of each flow layer is lower triangular, its determinant is simply the product of its diagonal entries [Huang et al., 2018]. This gives rise to the autoregressive flow formulation: given an ordering π of the d variables in the data vector x , the j th component of the flow T has the form: x j = τ j ( z j ; c j ( x <π ( j ) )) where each τ j is an invertible transformer and each c j is a conditioner that only depends on the variables that come before x j in the ordering π . As a result, the map components define an autoregressive model that factors the density over a random variable x as: p ( x ) = ∏ d j =1 p ( x j | x <j ) where x <j = ( x 1 , . . . , x j -1 ) . autoregressive . Under mild conditions, any arbitrary distribution p x can be transformed into a base distribution with a lower triangular Jacobian [Rezende and Mohamed, 2015]. That is, autoregressive flows are arbitrarily expressive given the target distribution. One common choice of invertible functions for the transformer are monotonic neural networks [Wehenkel and Louppe, 2019].

Density Estimation with Continuous Normalizing Flows: Continuous normalizing flows (CNFs) [Chen et al., 2018, Grathwohl et al., 2018] represent the transformation T as the flow map solving the differential equation ∂ z ( t ) ∂t = f ( z ( t ) , t ; θ ) . Given the initial condition x = z ( t 1 ) ∼ p x , we can integrate f backwards in time from t 1 to t 0 to obtain z = z ( t 0 ) ∼ p z , or vice versa. In order to learn the CNF, Chen et al. [2018] computes the change in log-density under the transformation using the instantaneous change of variables formula, which is defined by the differential equation ∂ log p ( z ( t )) ∂t = -Tr ( ∂f ∂ z ( t ) ) . This expression is used to compute the log-likelihood of a target sample as log p ( z ( t 1 )) = log p ( z ( t 0 )) -∫ t 1 t 0 Tr ( ∂f ∂ z ( t ) ) dt . We refer the reader to Chen et al. [2018] for the process of back-propagating through the objective using the adjoint sensitivity method as well as a discussion on the existence and uniqueness of a solution for the ODE flow map.

Causal Inference with Autoregressive Flows: Modelling causal relationships is crucial for enabling effective decision-making in various fields [Pearl, 2009]. A structural equation model (SEM) parameterizes the process that generates observed data, allowing us to reason about interventions and counterfactuals. Given random variables x = ( x 1 , ..., x d ) ∈ R d with joint distribution P x , the associated SEM consists of d structural equations of the form x j = f j ( pa j , u j ) , where u represents mutually independent latent variables and pa j denotes the direct causal parents of variable x j . Each SEM also corresponds to a directed acyclic graph (DAG), with a causal ordering π defined by the variables' dependencies.

Khemakhem et al. [2021] showcased the intrinsic connection between SEMs and autoregressive flows. The authors demonstrated that affine autoregressive flows with a fixed ordering of variables can be used to parameterize SEMs under the framework of causal autoregressive flows (CAREFL). When the causal ordering of variables is given, CAREFL outperforms other baselines in interventional tasks and generates accurate counterfactual samples. However, one significant limitation of CAREFL is that there is no guarantee that the autoregressive structure corresponds to the true dependencies in the causal graph beyond pairwise examples. In Section 3.3, we leverage StrAF to incorporate additional independence structure, enhancing its faithfulness to the causal DAG.

3 Methodology

3.1 Structured Neural Networks

As neural networks are universal function approximators [Hornik et al., 1989], injecting structure into multilayer perceptrons (MLPs) naturally allows us to approximate complex functions with certain invariances. For a neural network that approximates an arbitrary function f : R m → R n , one type of invariance we could consider is when a specific output ˆ x j is independent from a given input x i , i.e.: ˆ x j ⊥ x i means ∂ ˆ x j ∂x i = 0 . As a motivating example, we focus on the probabilistic density estimation problem framed as learning maps from one probability distribution to another. In this setting, it is imperative that we are able to learn structured functions between two high-dimensional spaces. Therefore in this paper, we seek to efficiently encode such invariances via weight masking.

For data x = ( x 1 , ..., x d ) , we use the lower-triangular adjacency matrix A ∈ { 0 , 1 } d × d to represent the underlying variable dependence structure. In other words, A ij = 0 for j < i if and only if x i ⊥ x j | x { 1 ,...,i }\ j and A ij = 1 otherwise. This matrix encodes the same information as a Bayesian network DAG of the variables. In the fully autoregressive case, matrix A is a dense lower triangular matrix with all ones under the diagonal, which is the only case addressed in Germain et al. [2015]. Their proposed MADE algorithm (Appendix A.1) only encodes the structure of dense adjacency matrices, and cannot incorporate additional conditional independencies. Further, the non-deterministic version of the algorithm can introduce unwanted independence statements by chance, as discussed in Proposition 1.

We improve upon the idea of masked autoregressive neural networks to directly encode the independence structure represented by an adjacency matrix A that is lower triangular but also has added sparsity. We observe that a masked neural network satisfies the structural constraints prescribed in A if the product of the masks for each hidden layer has the same locations of zero and non-zero entries as A . Therefore, given the conditional independence structure of the underlying data generating process, we can encode structure into an autoregressive neural network by factoring the adjacency matrix into binary mask matrices for each hidden layer .

More specifically, given an adjacency matrix A ∈ { 0 , 1 } d × d and a neural network with L hidden layers, each with h 1 , h 2 , ..., h L hidden units ( ≥ d ), we seek mask matrices M 1 ∈ { 0 , 1 } h 1 × d , M 2 ∈ { 0 , 1 } h 2 × h 1 , . . . , M L ∈ { 0 , 1 } d × h L such that A ′ ∼ A, where A ′ := M L · · · M 2 · M 1 . We use A ′ ∼ A to denote that matrices A ′ and A share the same sparsity pattern, i.e.: exact same locations of zeros and non-zeros. Note that here A is a binary matrix and A ′ is an integer matrix. We then mask the neural network's hidden layers using M 1 , M 2 , ..., M L as per (1) to obtain a Structured Neural Network (StrNN) , which respects the prescribed independence constraints. The value of each entry A ′ ij thus corresponds to the number connections flowing from input x j to output ˆ x i in the StrNN.

Finding the optimal solution to this problem is NP-hard since binary matrix factorization can be reduced to the biclique covering problem [Miettinen and Neumann, 2020, Ravanbakhsh et al., 2016, Orlin, 1977]. Furthermore, most existing works focus on deconstructing a given matrix A into low-rank factors while minimizing (but not eliminating) reconstruction error [Dan et al., 2015, Fomin et al., 2020]. In our application, any non-zero reconstruction error breaks the independence structure we want to enforce in our masked neural network. This puts existing algorithms for low-rank binary matrix factorization outside the scope of our paper. We instead consider the problem of finding factors that reproduce the adjacency matrix exactly, which is always possible when hidden layer dimensions are greater than the input and output dimension.

Optimization Objectives. Identifiability remains an issue even when we eliminate reconstruction error. Given an adjacency matrix A , there can be multiple solutions for factoring A into per-layer masks that satisfy the constraints, especially if the dimensions of the hidden layers are much larger than d . Since the masks dictate which connections remain in the neural network, the chosen mask factorization algorithm directly impacts the neural network architecture. Hence, it is natural to explicitly specify a relevant objective to the neural network's approximation error during the matrix factorization step, such as the test log-likelihood.

Given that the approximation error is inaccessible when selecting the architecture, we were inspired by the Lottery Ticket Hypothesis [Frankle and Carbin, 2018] and other pruning strategies [Srivastava et al., 2014, Gal and Ghahramani, 2016] that identify a subset of valuable model connections. Our hypothesis is that given the same data and prior knowledge on independence structure, the masked

neural network with more connections is more expressive, and will thus be able to learn the data better and/or more quickly. To find such models, we consider two objectives: Equation (2) that maximizes the number of connections in the neural network while respecting the conditional independence statements dictated by the adjacency matrix, and Equation (3) that maximizes connections while penalizing any pair of variables from having too many connections at the cost of the others. That is,

where var ( A ′ ) is the variance across all entries in A ′ . While we focus on these two objectives, future work will find optimal architectures by identifying other objectives to improve approximation error.

Factorization Algorithms. The maximization of the aforementioned mask factorization is an intractable optimization problem. We therefore develop approximate algorithms to solve them. Our strategy for optimization involves recursively factorizing the mask matrix layer by layer. Given A ∈ { 0 , 1 } d × d , we run Algorithm 1 once to find A 1 ∈ { 0 , 1 } d × h 1 and M 1 ∈ { 0 , 1 } h 1 × d for a layer of width h 1 such that A 1 · M 1 ∼ A , where ∼ denotes that the matrices share the same sparsity. For the next layer with h 2 hidden units, we use A 1 in place of A to find A 2 ∈ { 0 , 1 } d × h 2 and M 2 ∈ { 0 , 1 } h 2 × h 1 such that A 2 · M 2 ∼ A 1 . We repeat until we have found all the masks.

For each objective, we can obtain per-layer exact solutions using integer programming. While the Gurobi optimizer [Gurobi Optimization, LLC, 2023] can be used for small d , this approach was found to be too computationally expensive for d greater than 20, which is a severe limitation for real-world datasets and models. We hereby propose a greedy algorithm (shown in Algorithm 1 ) that approximates the solution to the maximum connections objective in Equation (2). For each layer, the algorithm first replicates the structure given in the adjacency matrix A by copying its rows into the first mask. It then maximizes the number of neural network connections by filling in the second mask with as many ones as possible while respecting the sparsity in A . See Appendix A.2 for a visual explanation of the algorithm. For a network with d -dimensional inputs and outputs and one hidden layer with h units, this algorithm runs in O ( dh ) time. Scaling up to L layers, where each hidden layer commonly contains O ( d ) units, the overall runtime is O ( d 2 L ) , which is much more efficient than the integer programming solutions. From our experiments, the greedy algorithm executes nearly instantaneously for dimensions in the thousands.

In Appendix A, we include detailed results from investigating the link between the neural network's generalization performance and the choice of mask factorization algorithm. We observe that while the exact solution to objective (2) achieves a higher objective value than the greedy approach, it has no clear advantage in density estimation performance. Moreover, we found that models trained with the two objectives, (2) and (3), provide similar performance. However, some datasets might be more sensitive to the exact objective. For example, problems with anisotropic nonGaussian structure may require neural architectures with more expressivity in some variables that may be favored with certain objectives. While we adopted objective (2) and the efficient greedy algorithm in the remainder of our experiments, designing and compar-

ing different factorization objectives is an important direction for future work.

3.2 Structured Neural Networks for Normalizing Flows

For the general real-valued probabilistic density estimation task, we highlight the use of StrNN in two popular normalizing flow frameworks, autoregressive flows and continuous-time flows. In the autoregressive flows setting, recall that the j th component of each flow layer is parameterized as x j = τ j ( z j ; c j ( x <π ( j ) )) , where the conditioner c j dictates which inputs the latent variable can depend on. We use the StrNN as the conditioner, combined with any valid invertible transformer τ , to form the Structured Autoregressive Flow (StrAF) . The StrNN conditioner ensures each transformed latent

variable is only conditionally dependent on the subset of variables defined by a prescribed adjacency matrix. A single step from StrAF can be represented as: z = StrAF ( x ) = τ ( x , StrNN ( x , A )) where x ∈ R d and A ∈ { 0 , 1 } d × d is an adjacency matrix. Normalizing flow steps are typically composed to improve the expressiveness of the overall flow. To ensure that the prescribed adjacency structure is respected throughout all flow steps, we simply avoid the common practice of variable order permutation between flow steps. We visualize this in Figure 2. This choice does not hurt performance as long as a sufficiently expressive transformer is used. In our experiments, we apply the unconstrained monotonic neural network (UMNN) described in Wehenkel and Louppe [2019] as a transformer.

We also integrate the StrNN into continuous normalizing flows (CNFs), where the neural network that parameterizes f in the differential equation ∂ z ( t ) ∂t = f ( z ( t ) , t ) is replaced by a StrNN to obtain ∂ z ( t ) ∂t = StrNN ( z ( t ) , t, A ) . This plug-in replacement allows us to inject structure into the function describing the continuous dynamics of the CNF without modifying other aspects of the CNF. We refer to this architecture as the Structured CNF (StrCNF) . The StrCNF uses the trace estimator described in FFJORD [Grathwohl et al., 2018] to evaluate the objective function for learning the flow.

While injecting structure, both StrAF and StrCNF inherit the efficiency of StrNN due to our choice of weight masking. Specifically, the output of the StrNN can be computed with a single forward pass through the network. In comparison, input masking

Structured Autoregressive Flow (StrAF)

NF Step

UMNN

Connectivity

after 1 step

Final

Connectivity

True

Adjacency

Connectivity

after 2 steps

NF Step

UMNN

approaches such as Wehenkel and Louppe [2021] must perform d forward passes to compute the output for a single datum. This not only prevents efficient application of input masking to high dimensional data, but also is a barrier to integrating the method with certain architectures. For example, the CNF already requires many neural network evaluations to numerically solve the ODE defining the flow map, so making d passes per evaluation is particularly inefficient.

3.3 Structured Causal Autoregressive Flow

Beyond density estimation, StrAF allows us to build autoregressive flows that faithfully represent variable dependencies defined by a causal DAG. In contrast, CAREFL [Khemakhem et al., 2021] only maintains the autoregressive order, which is insufficient for data with more than two variables where the true structure must be characterized by a full adjacency matrix. Building on CAREFL, we also assume that the flow T takes on the following affine functional forms for observed data x :

where the functions s j and t j are both conditioners that control the dependencies on the variables preceding x j . It is crucial for the autoregressive or graphical structures to be maintained across all sub-transformations T 1 , ..., T K for the flow T = T 1 ◦ · · · ◦ T K .

Assuming the true causal topology has been given either from domain experts or oracle discovery algorithms, StrAF can directly impose the given topological constraints by adding an additional masking step based on the adjacency matrix, as shown in section 3.1. This ensures that the dependencies of the flows match the known causal structure, which leads to more accurate inference predictions.

4 Related Works

Given a Bayesian network adjacency matrix, Wehenkel and Louppe [2021] introduced graphical conditioners to the autoregressive flows architecture through input masking. They demonstrated that unifying normalizing flows with Bayesian networks showed promise in injecting domain knowledge while promoting interpretability, as even single-step graphical flows yielded competitive results in density estimation. Our work follows the same idea of introducing prior domain knowledge into

autoregressive flows, but we instead use a masking scheme similar to methods in Germain et al. [2015]. We again note that input masking scales poorly with data dimension d , as d forward passes are required to obtain the output for a single datum.

Silvestri et al. [2021] proposed embedded-model flows, which alternates between traditional normalizing flows layers and gated structured layers that a) encode parent nodes based on the graphical model, and b) include a trainable parameter that determines how strongly the current node depends on its parent nodes, which alleviates error when the assumed graphical model is not entirely correct. In comparison, our work encodes conditional independence more directly in the masking step, improving accuracy when the assumption in the probabilistic graphical structure is strong.

Mouton and Kroon [2022] applied a similar idea to residual flows by masking the residual blocks' weight matrices prior to the spectral normalization step according to the assumed Bayesian network. Similarly, Weilbach et al. [2020] introduced graphical structure to continuous normalizing flows by masking the weight matrices in the neural network that is used to parameterize the time derivative of the flow map. However, their method is difficult to apply to neural networks with more than a single hidden layer. In comparison, our factorization algorithm more naturally permits the use of multi-layered neural networks when representing CNF dynamics. The Zuko software package [Rozet et al., 2023] implements various types of normalizing flows, including MAFs that also rely on weight masking, but they only provide one possible algorithm to enforce autoregressive structure given a specific variable order. In comparison, our approach permits explicit optimization of different objectives during the adjacency matrix factorization step and we investigate the efficacy of factorization schemes and resulting neural architectures in our work. For completeness, we include pseudocode of their algorithm and comparisons to our own algorithms in Appendix A .

Flows have garnered increasing interest in the context of causal inference, with applications spanning various problem domains. Ilse et al. [2021] parameterized causal model with normalizing flows in the general continuous setting to learn from combined observational and interventional data. Melnychuk et al. [2022] used flows as a parametric method for estimating the density of potential outcomes from observational data. Flows have also been employed in causal discovery [Brouillard et al., 2020, Khemakhem et al., 2021] as well as in various causal applications [Ding et al., 2023, Wang et al., 2021]. In particular, Balgi et al. [2022a] also considered embedding the true causal DAG in flows for interventional and counterfactual inference, but they do so via the framework of Graphical Normalizing Flows [Wehenkel and Louppe, 2021]. Balgi et al. [2022b] used CAREFL on a real-world social sciences dataset leveraging a theorized Bayesian network.

5 Experiments

To demonstrate the efficacy of encoding structure into the learning process, we show that using StrNN its flow integrations to enforce a prescribed adjacency structure improves performance on density estimation and sample generation tasks. We experiment on both synthetic data generated from known structure equations and MNIST image data. Details on the data generation process for all synthetic experiments can be found in Appendix C. We also apply StrAF in the context of causal inference and demonstrate that the additional graphical structure introduced by StrAF leads to more accurate interventional and counterfactual predictions. The code to reproduce these experiments is available at https://github.com/rgklab/StructuredNNs .

5.1 Density Estimation on Binary Data

For binary density estimation, we compare StrNN against the fully autoregressive MADE baseline.

Synthetic Tabular Data We generate binary tabular data through structural equations from known Bayesian networks. The results are shown in Figure 3 (left). We find that StrNN performs better than MADE, especially in the low data regime, as demonstrated on the left hand side of each chart.

MNIST Image Data To study the effect of structure in image modeling, we use the binarized MNIST dataset considered in Germain et al. [2015], Salakhutdinov and Murray [2008]. Germain et al. [2015] treated each 28-by-28-pixel image as a 784-dimensional data vector with full autoregressive dependence. Since we do not know the ground truth structure, we use StrNN to model a local autoregressive dependence on a square of a pixels determined by the hyper-parameter nbr_size . By changing the hyperparameter we can increase the context window used to model each pixel.

Synthetic:Binaryevery_other

Synthetic:Binaryrandom_sparse

BinaryMNiSTLabel2

29

250

MADE

MADE

MADE

150

MADE

150

LoSS

225

28

140

140

Test

200

27

130

130

175

120

120

26

110

150

110

2000

40006000800010000

0

2000

4000

500

1000

1500

2000

500

1000

1500

2000

SampleSize

SampleSize

SampleSize

SampleSize

StrNN is equivalent to MADE for this experiment when we set nbr_size=28 . We first find the best nbr_size for each label via grid search. For labels 0 and 2, the optimal nbr_size is 10, and per-label density estimation results can be found in Figure 3 (right). StrNN outperforms MADE for both labels, with the advantage more significant when sample size is very small. Samples of handwritten digits generated from both StrNN and MADE can be found in Appendix B.

5.2 Density Estimation on Gaussian Data

We run experiments to compare the performances of StrNN and MADE on synthetic Gaussian data generated from known structure equation models where each x i is Gaussian. We plot the results in Figure 4. StrNN achieves lower test loss than MADE on average, although the error bars are not necessarily disjoint. When the sample size is low, however, StrNN significantly outperforms MADE, similar to the binary case. In conclusion, across all binary and Gaussian experiments, encoding structure makes StrNN significantly more accurate at density estimation than the fully autoregressive MADE baseline.

Synthetic:Gaussianprev_2

Synthetic:Gaussianrandom_sparse

MADE

4.5

MADE

-13

LoSS

-14

4.0

15

3.5

-16

3.0

17

2.5

0

2000

4000

1000

2000

3000

4000

SampleSize

SampleSize

5.3 Density Estimation with Structured Normalizing Flows

We evaluate StrAF on density estimation against several baselines. We draw 1000 samples from a 15 dimensional tri-modal and non-linear synthetic dataset for experimental evaluation. Data generation is further described in Appendix C.4.

Experimental Setup We select the fully autoregressive flow (ARF) and the Graphical Normalizing Flow (GNF) [Wehenkel and Louppe, 2021] as the most relevant discrete flow baselines for comparison. We use the official GNF code repository during evaluation, but note that it contains design decisions that harm sample quality (see Appendix E.3). We also evaluate against FFJORD [Grathwohl et al., 2018] and Weilbach et al. [2020] as baselines for StrCNF. While other structured flows exist and have been examined in Section 4, they do not represent variables using an autoregressive structure and hence are less comparable. Where applicable, models were provided the true adjacency matrix in their conditioners. All discrete flow models use a UMNN [Wehenkel and Louppe, 2019] transformer and we grid-search other hyperparameters as described in Appendix E.2. We evaluate density estimation performance using the negative log-likelihood (NLL). After fixing the hyperparameters, we perform 8 randomly initialized training runs, then report the mean and 95% CI of the test NLL for these runs.

StrNN improves flow-based models for density estimation

Wereport results in Table 1 and observe several trends. For both discrete and continuous flows, the ability to incorporate structure yields performance benefits compared to the ARF and FFJORD baselines. In particular, the StrNN offers advantages in comparison to baseline approaches that can encode structure. For example, while the method proposed by Weilbach et al. [2020] allows a DAG structure to

(

X

1,

X

4)

ARF

GNF

FFJORD

Weilbach

StrAF

StrCNF

(

X

2,

X

12)

(

X

4,

X

10)

be injected into the baseline CNF, its inability to easily use a multilayered neural network to represent dynamics hinders performance as compared to the StrCNF. We observe that StrAF and GNF perform comparably, and the StrCNF outperforms all other models. We visualize the quality of samples generated by these flow models in Figure 5, and find that both StrAF and StrCNF yields samples that closely match the ground truth distribution.

5.4 Causal Inference with Structured Autoregressive Flows

We conduct synthetic experiments where we generate data according to a linear additive SEM. Details on data generation can be found in Appendix D.2. In our experiments involving 5- and 10-variable

SEMs, we compare StrAF against CAREFL, which only utilizes MADE as the conditioner and maintains autoregressive ordering of the variables. This comparison highlights the additional benefits of enforcing the generative model to be faithful to the causal graph, which is a feature unique to StrAF. Furthermore, unlike the previous work conducted by Khemakhem et al. [2021] that evaluates causal queries on individual variables alone, we propose a comprehensive evaluation metric called total mean squared error (MSE) for these two causal tasks as outlined below. We report the mean errors along with the standard deviations from multiple training runs with different datasets. The formulations of the metrics can be found in Appendix D.3. In addition, Appendix D.1 provides detailed algorithms and related discussions on generating interventional samples and computing counterfactuals with flows.

500

1000

1500

2000

Sample size

0.0

0.5

1.0

1.5

Total I-MSE

5-variable SEM

CAREFL

StrAF

500

1000

1500

2000

Sample size

0.0

0.5

1.0

1.5

10-variable SEM

CAREFL

StrAF

(a) Interventions

500

1000

1500

2000

Sample size

0

2

4

6

Total C-MSE

5-variable SEM

CAREFL

StrAF

500

1000

1500

2000

Sample size

0

5

10

15

10-variable SEM

CAREFL

StrAF

(b) Counterfactuals

Interventions. We perform interventions on each variable x j within the SEM and draw samples from the intervened causal system. Interventions are performed by setting the intervened value α using one of eight integers perturbed from the mean of the intervened variable. We compute the expectations E [ x i | do ( x j = α )] for variable x i , excluding the intervened variable x j itself and its

preceding variables as they remain unaffected. To estimate these expectations, we use 1000 samples and calculate the squared error for the predictions. The resulting error is then averaged over the set of intervened values, intervened variables, and the corresponding variables with computed expectations under interventions. We present the plot of this aggregated metric, referred to as the total intervention mean squared error (total I-MSE), for the two SEMs in Figure 6a. Notably, StrAF consistently outperforms CAREFL across all training dataset sizes, highlighting the effectiveness of the additional graph structure in improving StrAF's inference of the interventional distribution.

Counterfactuals. In a similar setting as the intervention experiments, we evaluate StrAF and CAREFL on their ability to compute accurate counterfactuals conditioned on observed data. Counterfactual inference tackles what-if scenarios: determining the value of variable x i if variable x j had taken a different value α . Unlike interventions, counterfactual queries involve deriving the latent variables z given the observed data x = x obs, rather than sampling new z . We generate 1000 observations using the synthetic SEM and derive counterfactual values x by posing queries with varying α values for each variable x j . Similar to interventions, we compute the squared error and average over the 1000 observed samples and all possible combinations of counterfactual queries for each observed sample, and we refer to this metric as the total counterfactual mean squared error (total C-MSE). Figure 6b illustrates that, for both SEMs, StrAF outperforms CAREFL in making more accurate counterfactual predictions. Moreover, StrAF demonstrates consistent performance even in scenarios with limited available samples.

6 Conclusions and Limitations

We introduce the Structured Neural Network (StrNN) , which enables us to encode functional invariances in arbitrary neural networks via weight masking during learning. For density estimation tasks where the true dependencies are expressed via Bayesian networks (or adjacency matrices), we show that StrNN outperforms a fully autoregressive MADE model on synthetic and MNIST data. We integrate StrNN in autoregressive and continuous flow models to improve both density estimation and sample quality. Finally, we show that our structured autoregressive flow-based causal model outperforms existing baselines on causal inference. We address some limitations and directions for future work below.

Access to true adjacency structure. Weassume access to the true conditional independence structure. While this information is available in many contexts, such as from domain experts, there is also a wide literature on learning the conditional independencies directly from data [Drton and Maathuis, 2017]. One prominent example is the NO-TEARS algorithm [Zheng et al., 2018]. Wehenkel and Louppe [2021] shows that integrating NO-TEARS with an autoregressive flow can improve density estimation when a ground truth adjacency is unknown. Any adjacency matrix learned from data can also be integrated in StrAF. In our work, we have found two existing causal structure discovery libraries that are relatively comprehensive and easy to use: Kalainathan et al. [2020] and Zheng et al. [2023]. Future work will use StrNN to directly learn structure from data, providing a full pipeline from structure discovery to density estimation and sample generation.

StrNN optimization objectives. The mask factorization algorithm used by StrNN can maximize different objectives while ensuring the matrices satisfy a sparsity constraint. In Section 3.1, we proposed two such objectives and in Appendix A we demonstrated that they can impact model generalization. StrNN provides a framework with which it is possible to explore other objectives to impose desirable properties on neural network architectures. Moreover, investigating the effect of sparse structure on faster and easier training is a valuable direction [Frankle and Carbin, 2018]. For example, one approach would be to leverage connections between dropout and the lottery ticket hypothesis to randomly introduce sparsity into a neural network using StrNN weight masking.

Acknowledgments and Disclosure of Funding

We thank David Duvenaud, Tom Ginsberg, Vahid Balazadeh-Meresht, Phil Fradkin, and Michael Cooper for insightful discussions and draft reviewing. We thank anonymous reviewers for feedback that has greatly improved the work. This research was supported by an NSERC Discovery Award RGPIN-2022-04546, and a Canada CIFAR AI Chair. AC is supported by a DeepMind Fellowship and RS is supported by the Ontario Graduate Scholarship and the Ontario Institute for Cancer Research. Resources used in preparing this manuscript were provided in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.

References

Jungseul Ok, Alexandre Proutiere, and Damianos Tranos. Exploration in structured reinforcement learning. Advances in Neural Information Processing Systems , 31, 2018.

Zeshan M Hussain, Rahul G Krishnan, and David Sontag. Neural pharmacodynamic state space modeling. In International Conference on Machine Learning , pages 4500-4510. PMLR, 2021.

Limeng Cui, Haeseung Seo, Maryam Tabar, Fenglong Ma, Suhang Wang, and Dongwon Lee. Deterrent: Knowledge guided graph attention network for detecting healthcare misinformation. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining , pages 492-502, 2020.

Ali Hossein Foomani Gharari, Michael Cooper, Russell Greiner, and Rahul G Krishnan. Copula-based deep survival models for dependent censoring. In Uncertainty in Artificial Intelligence , pages 669-680. PMLR, 2023.

Sebastian Curi, Silvan Melchior, Felix Berkenkamp, and Andreas Krause. Structured variational inference in partially observable unstable gaussian process state space models. In Learning for Dynamics and Control , pages 147-157. PMLR, 2020.

Vahid Balazadeh, Vasilis Syrgkanis, and Rahul G Krishnan. Partial identification of treatment effects with implicit generative models. arXiv preprint arXiv:2210.08139 , 2022.

Esteban G Tabak and Cristina V Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics , 66(2):145-164, 2013.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference on machine learning , pages 1530-1538. PMLR, 2015.

Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE: Masked autoencoder for distribution estimation. In International Conference on Machine Learning , pages 881-889. PMLR, 2015.

Judea Pearl. Bayesian networks. 2011.

George Papamakarios, Theo Pavlakou, and Iain Murray. Masked autoregressive flow for density estimation. Advances in neural information processing systems , 30, 2017.

Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In International Conference on Machine Learning , pages 2078-2087. PMLR, 2018.

Judea Pearl. Causality . Cambridge university press, 2009.

George Papamakarios, Eric T Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference. J. Mach. Learn. Res. , 22(57):1-64, 2021.

Youssef Marzouk, Tarek Moselhy, Matthew Parno, and Alessio Spantini. An introduction to sampling via measure transport. arXiv preprint arXiv:1602.05023 , 2016.

Antoine Wehenkel and Gilles Louppe. Unconstrained monotonic neural networks. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems , volume 32. Curran Associates, Inc., 2019. URL https://proceedings.neurips.cc/paper_files/paper/2019/file/ 2a084e55c87b1ebcdaad1f62fdbbac8e-Paper.pdf .

Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems , 31, 2018.

Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367 , 2018.

Ilyes Khemakhem, Ricardo Monti, Robert Leech, and Aapo Hyvarinen. Causal autoregressive flows. In International conference on artificial intelligence and statistics , pages 3520-3528. PMLR, 2021.

Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are universal approximators. Neural networks , 2(5):359-366, 1989.

Pauli Miettinen and Stefan Neumann. Recent developments in boolean matrix factorization. arXiv preprint arXiv:2012.03127 , 2020.

Siamak Ravanbakhsh, Barnabás Póczos, and Russell Greiner. Boolean matrix factorization and noisy completion via message passing. In International Conference on Machine Learning , pages 945-954. PMLR, 2016.

James Orlin. Contentment in graph theory: Covering graphs with cliques. Indagationes Mathematicae (Proceedings) , 80(5):406-424, 1977. ISSN 1385-7258. doi: https://doi.org/10.1016/ 1385-7258(77)90055-5. URL https://www.sciencedirect.com/science/article/pii/ 1385725877900555 .

Chen Dan, Kristoffer Arnsfelt Hansen, He Jiang, Liwei Wang, and Yuchen Zhou. Low rank approximation of binary matrices: Column subset selection and generalizations. arXiv preprint arXiv:1511.01699 , 2015.

Fedor V Fomin, Petr A Golovach, and Fahad Panolan. Parameterized low-rank binary matrix approximation. Data Mining and Knowledge Discovery , 34:478-532, 2020.

Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv:1803.03635 , 2018.

Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: a simple way to prevent neural networks from overfitting. The journal of machine learning research , 15(1):1929-1958, 2014.

Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In Maria Florina Balcan and Kilian Q. Weinberger, editors, Proceedings of The 33rd International Conference on Machine Learning , volume 48 of Proceedings of Machine Learning Research , pages 1050-1059, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https://proceedings.mlr.press/v48/gal16.html .

Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL https://www. gurobi.com .

Antoine Wehenkel and Gilles Louppe. Graphical normalizing flows. In International Conference on Artificial Intelligence and Statistics , pages 37-45. PMLR, 2021.

Gianluigi Silvestri, Emily Fertig, Dave Moore, and Luca Ambrogioni. Embedded-model flows: Combining the inductive biases of model-free deep learning and explicit probabilistic modeling. arXiv preprint arXiv:2110.06021 , 2021.

Jacobie Mouton and Steve Kroon. Graphical residual flows. arXiv preprint arXiv:2204.11846 , 2022.

Christian Weilbach, Boyan Beronov, Frank Wood, and William Harvey. Structured conditional continuous normalizing flows for efficient amortized inference in graphical models. In Silvia Chiappa and Roberto Calandra, editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics , volume 108 of Proceedings of Machine Learning Research , pages 4441-4451. PMLR, 26-28 Aug 2020. URL https://proceedings.mlr.press/v108/ weilbach20a.html .

Francois Rozet, Felix Divo, and Simon Schnake. Zuko: Normalizing Flows in PyTorch. https://github.com/francois-rozet/zuko, 2023. Version 0.3.1.

Maximilian Ilse, Patrick Forré, Max Welling, and Joris M Mooij. Combining interventional and observational data using causal reductions. arXiv preprint arXiv:2103.04786 , 2021.

Valentyn Melnychuk, Dennis Frauen, and Stefan Feuerriegel. Normalizing flows for interventional density estimation. arXiv preprint arXiv:2209.06203 , 2022.

Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems , 33:21865-21877, 2020.

Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Causalaf: Causal autoregressive flow for safety-critical driving scenario generation. In Conference on Robot Learning , pages 812-823. PMLR, 2023.

Rongguang Wang, Pratik Chaudhari, and Christos Davatzikos. Harmonization with flow-based causal inference. In Medical Image Computing and Computer Assisted Intervention-MICCAI 2021: 24th International Conference, Strasbourg, France, September 27-October 1, 2021, Proceedings, Part III 24 , pages 181-190. Springer, 2021.

Sourabh Balgi, Jose M Pena, and Adel Daoud. Personalized public policy analysis in social sciences using causal-graphical normalizing flows. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 11810-11818, 2022a.

Sourabh Balgi, Jose M Peña, and Adel Daoud. Counterfactual analysis of the impact of the imf program on child poverty in the global-south region using causal-graphical normalizing flows. arXiv preprint arXiv:2202.09391 , 2022b.

Ruslan Salakhutdinov and Iain Murray. On the quantitative analysis of deep belief networks. In International Conference on Machine Learning , 2008.

Mathias Drton and Marloes H Maathuis. Structure learning in graphical modeling. Annual Review of Statistics and Its Application , 4:365-393, 2017.

Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing. Dags with no tears: Continuous optimization for structure learning. Advances in neural information processing systems , 31, 2018.

Diviyan Kalainathan, Olivier Goudet, and Ritik Dutta. Causal discovery toolbox: Uncovering causal relationships in python. The Journal of Machine Learning Research , 21(1):1406-1410, 2020.

Yujia Zheng, Biwei Huang, Wei Chen, Joseph Ramsey, Mingming Gong, Ruichu Cai, Shohei Shimizu, Peter Spirtes, and Kun Zhang. Causal-learn: Causal discovery in python. arXiv preprint arXiv:2307.16405 , 2023.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work studies the impact of imposing known conditional independence structure in fully-connected neural network architectures, notably in the setting of autoregressive normalizing flows. The independence is imposed by masking the weight matrices in the linear layers, similar to the MADE approach. The masks are determined by a factorization of the known adjacency matrix that (approximately) maximizes the number of connections (paths) between inputs and outputs. Through several experiments, the authors show generalization improvement over baselines.\n\n**Edit** Most of my concerns were addressed during the rebuttal/discussion period. I am therefore upgrading my score from 4 to 6.\n\nSTRENGTHS:\n* The manuscript is well written, with standard notations.\n* The contributions are clearly defined and the assumptions (known adjacency) are explicit.\n* The proposed approach to impose conditional independences is sound.\n* The claims are supported by the experiments, notably concerning improved generalization.\n\nWEAKNESSES:\nThe idea of imposing prior independence knowledge into autoregressive flows was first introduced by Wehenkel and Louppe (2021), cited as [25] in this manuscript. As the authors mention (lines 246-247), StrAF only differs from GNF in the approach to impose the independences, but is conceptually identical. Actually, Wehenkel and Louppe (2021) already propose the approach of the present work:\n\n    > An alternative approach would be to use masking scheme similar to what is done by Germain et al. (2015) in MADE as suggested by Lachapelle et al. (2019).\n    \nIn addition, the official [UMNN repository](https://github.com/AWehenkel/UMNN), also cited in this work, links to the normalizing flow library [Zuko](https://github.com/francois-rozet/zuko), from the same lab. The latter library implements autoregressive flow conditioners as masked multi-layer perceptrons for which the masks are a factorization of the adjacency matrix between the inputs and outputs. The similarities with the proposed StrNN are too strong to be left unaddressed.\n\nQUESTIONS:\n* Algorithm 1: It is not clear to me how the factorization algorithm is applied when the StrNN has more than one hidden layer.\n* Section 5.1: Are the same number of layers/neurons used for StrNN and MADE in this experiment?\n* Line 317: \"As GNF permutes variables between flow steps\". I was not able to find a mention of this in [25].\n* Figure 5/Table 1: I don't understand how \"ARF-10\" and \"GNF-10\" can be so much worse than \"GNF-1\", as they are strictly more expressive. Is it an overfitting issue? Or maybe an invertibility issue? UMNN is not always numerically invertible. What about \"ARF-1\"?\n* Table 1: The authors make a distinction between \"density estimation\" and \"sample quality\", which does not make sense to me. If a flow perfectly estimates the density, it necessarily generates probable samples, unless the invertibility is not guaranteed.\n* Why not studying the use of StrNN in other settings than density evaluation, such as physics-informed machine learning, where it is common to infuse prior knowledge in the structure of the neural networks?\n\nLIMITATIONS:\n* It is never mentioned that a StrNN is a (pruned) fully-connected network with element-wise activation functions. The approach does not apply to convolutional, attention-based or recurrent networks, and does not support skip/residual connections or normalization layers.\n* This is not a limitation of this work, but one should be careful not to confuse Bayesian networks and causal graphs. A Bayesian network (or its adjacency matrix) over variables merely indicates independencies between the variables but in no way causalities.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a novel method for constructing structured neural networks (strNN) that are able to respect causal independencies between variables. Formal constraints for weight mask creation are discussed and evaluated empirically for an exact method and a greedy algorithm. The conditioned strNN are furthermore leveraged for conditioning autoregressive flow models. Practical application successfully is demonstrated over multiple synthetic datasets with comparisons between several baseline models, with and without the use of adjacency information. To the best of my knowledge related work is discussed sufficiently in the context of causal density estimation.\n\nSTRENGTHS:\n1. The authors propose a novel method for constructing structured models via weight masking that integrates seamlessly with existing neural architectures while respecting the strict independence assumptions of causal models. Preconditions and assumptions for application of the approach, specifically knowledge about the causal graph structure, are clearly stated.\n\n2. To tackle the infeasibility of exact mask creation on larger graphs a greedy algorithm is proposed and its practical application is demonstrated.\n\n3. The authors additionally include experiments on binary MNIST data, for which the underlying causal structure is unknown. By imposing a causal graph structure which promotes the usage of spatially local information, the authors improve performance for non-synthetic data over baselines.\n\nWEAKNESSES:\nThe example shown in Figure 1 decomposes the network into two separate networks. However, constructing the displayed a network would not require a complicated mask decomposition, but could be trivially solved by constructing two independent networks with constrained layer width. Only by inspecting the example provided in the appendix it is revealed that the presence of split-structures leads to shared weights between the outputs.\n\nQUESTIONS:\n1. As causal mechanisms are often assumed to be independent in causal literature, I would like to ask the authors about the benefits or downsides of allowing for such shared weights within the network.\n\n2. Furthermore, I would like to ask the authors to discuss possible simplifications for specific causal structures, e.g. in the case of independent causal mechanisms as seen in Figure 1.\n\nOverall the idea is pretty good with a clever way of enforcing the causal independencies. However, all of this is assuming that the networks can leverage shared information between different mechanisms. If that is not the case then you could just train an independent density estimator for every single edge and (from a purely causal perspective) the problem becomes trivial to solve.\n\nLIMITATIONS:\nNo concerns here", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors present a neural network architecture that can fulfill the bayesian DAG conditional independencies needed for normalized density estimation.\nIn this work, the authors start from a binary lower adjacency matrix that encodes the independencies of a bayesian network DAG. Then they introduce a factorization of the global adjacency matrix into L adjacency matrices, which can be used as masks on each layer. This construction allows neural networks to be trained for a normalized density estimation task.\nThe authors introduce an heuristic to exactly factorize the adjacency matrix for the different layers using two objective functions.\nWith this building block the authors proceed to create a normalizing flow architecture that respects the independency restrictions end-to-end.  The authors then compare their approach empirically on different task against MADE, a neural density estimator that allows general dependencies for a given random variable ordering.\n\nSTRENGTHS:\nI found the paper insightful, and the results show that the approach is beneficial. The paper is technically sound and easy to read. The results showing an improvement in data efficiency for a given negative likelihood are also very interesting.\n\nThe introduction of the normalizing flow approach and the comparison in the causal setting are also nice additions that can have impact in the broader community.\n\nThe experiment on the sample quality shows the benefit of restricting the dependencies in the network as it cuts paths for noise in other random variables (and feature transformations) to propagate through the network.\n\nWEAKNESSES:\nThe major weakness of this paper is the limited empirical section in comparison to other papers in this domain. This can cause readers to wonder if the benefits of the new approach as density estimators are not significant for other datasets.\nA broader comparison on other datasets would make the paper more robust. \n\nAlso and I'm considering this as a minor weakness in my review, is that the method although insightful does not provide a way to obtain the global adjacency matrix.\n\nQUESTIONS:\nAs I was reading the paper, my first thought would be that you would explore the possibility of discovering the dependencies for a given order. \nHere one can start with a dense adjacency matrix, train the network, clip nodes in the layers according to Lottery Ticket Hypothesis, and propagate the masks forward, i.e., multiplying all the masks to get the adjacency matrix. As you are clipping, the adjacency matrix is guaranteed to either be the same or introduce independencies. \nAt that point, you can even use your factorization algorithm again to obtain a network with more/other nodes active while still respecting the new independencies. \nI'm wondering if you explored similar ideas during your research?\n\nLIMITATIONS:\nThe main limitations of the approaches presented are inherited from the restrictions on ordered models, e.g., marginalization and map queries are intractable for the general case. This is not mentioned in the paper. \n\nThere are no potential negative societal impact to this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: SUMMARY:
This paper introduces structured neural networks such that the resulting neural network represents the factorization of a given Bayesian network. For doing so, each layer of the neural network is masked and the product of the masks of all layers must be the same as the adjacency matrix of the DAG representing the Bayesian network. 
With this construction, the represented conditional dependencies with structured neural networks will be consistent with the given Bayesian network. The paper proposes a simple greedy algorithm to find the masks. It also proposes using the structured neural network to construct the coupling layers for normalizing flow and claims that the resulting generative model is better for casual inference (intervention and counterfactuals) than the prior approach.

STRENGTHS:
The paper is well-written and the contribution towards causal inference is solid.

WEAKNESSES:
1- The structured neural network augments MADE with a better mask construction algorithm such that the factorization can be defined for any DAG structure. However, it is not a fundamentally different model.
2- The paper didn't propose any approach for learning the structure of DAG given the provided structured neural network parametrization.
3- MADE is not a strong density estimator and comparing only to MADE does not validate the strength of structured neural networks as density estimators.

QUESTIONS:
1) How GNF would compare to StrAF if it does permute the latent variables after the first step? 
2) During the comparison with CAREFL did you provide CAREFL with external DAG orders that StrAF uses? If not, the learned causal order by CAREFL may not exactly specify the underlying DAG, which may result in lower performance in causal inference.

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-27