Structured Neural Networks for Density Estimation and Causal Inference
Authors: Asic Q Chen
Ruian Shi
Xiang Gao
Ricardo Baptista
Rahul G Krishnan
Venue: NeurIPS 2023 poster (2023)
Keywords: generative models
density estimation
normalizing flows
binary matrix factorization
causal inference
TL;DR: We introduce an efficient way to inject structure into arbitrary neural networks based on weight masking, which has applications in density estimation, generative modelling, 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) 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
Asic Chen' Ruian Shi
Xiang Gao' Ricardo Baptista Rahul G. Krishnan 't
UUniversity of Toronto, Vector Institute {asicchen, ruiashi, xgao, rahulgk/0cs. toronto.... edu rsb8caltech. edu
CClifornia Institute of Technology
Abstract
Injecting structure into neural networks enables learning functions that satisfy invariances with respect to subsets of inputs. For instance, when learning gener- ative models using neural networks, it is advantageous to encode the conditional independence structure of observed variables, often in the form of Bayesian net- works. We propose the Structured Neural Network (StrNN), which injects structure through masking pathways
neural network. The masks are designed via 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. demonstrate 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 a conditioner, and (3) interventional and counterfactual analysis with StrAFs. Our work opens up avenues for learning neural networks that enable data-efficient generative modeling and the use of normalizing flows for causal effect estimation.
Introduction
p(X) Pb. (m:)ps (sagarr))o,, (sty (n))po
The incorporation of structure into machine learning models has been shown to provide benefits for model generalization, learning efficiency, and interpretability. The improvements particularly salient when learning from small amounts of data. This idea has found use in reinforcement learning [Ok al., 2018], computational healthcare [Hussain et al., 2021, Cui et al., 2020], sur- vival analysis [Gharari et al., 2023], time series analysis [Curi et al., 2020], and causal inference [Balazadeh et al., 2022].
25
No by
4: 4,
X
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 series of invertible transformations. Masked autoencoders (MADE) [Germain et al., 2015] model the joint distribution via an autoregressive factorization of
Weight
Welght
Weight Mask
Adjacency
Mask
Mask
Matrix
Figure 1: StrNN injects structure by masking weights neural network. Top: StrNN connections (green) compared to fully con- nected network (gray). Bottom: Binary factor- ization of adjacency matrix yields weight masks. Masked weights shown in gray.
Equaa Contribution 'Euual Senior Authorship
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

the random variables. The factorization is enforced by integrating structure in the neural network of the autoencoder. The MADE architecture zeros out weights 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 Bayesian network [Pearl, 201 1] 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 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 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 normal- izing 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 small number of samples.
In summary, the main contributions of this work are as follows:
We introduce StrNN, a veight-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.
2. We integrate StrNN into autoregressive and continuous normalizing flows for best-in-class perfor- mance in density estimation and sample generation tasks.
3. We apply StrAF for causal effect estimation and showcase its ability 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 = ...) MADE factorizes p(x) the product of the outputs of a neural network. Writing the j-th output as the conditional probability æÚ- := p(æ; = 1|x=;), the joint distribution can be rewritten exactly as the binary cross-entropy loss. As long as the neural network outputs are autoregressive in relation its inputs, we can minimize the cross-entropy loss for density estimation. To enforce the autoregressive property for neural network y = fz) with a single hidden layer and d inputs and outputs, MADE element-wiis multiplies the weight matrices W and V with binary masks M W and Mr:
h(z)= o((W @ M W )r + b), = f((V + M VM)) )1(z) + cc.
(I)
2

---

The autoregressive property is satisfied as long as the product of the masks, M E B 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 Rd - TR:2d (a11....TEd) - (ύ², (.1,,.... id, (og(()),... log (a)). The last mask must be duplicated to ensure Hy and o only depend on xzj. 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 distri- butions and have been applied in many scenarios [Papamakarios et al., 2021]. Given data x e R a normalizing flow T: I2 + Rd takes x to latent variables z E R' that are distributed according to a simple base distribution Pz. 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 soungeof--arribles formula: Px Px(x)= = P2 (T(x)) det JT(x)]. We can compose multiple diffeomor- phic transformations T to form the flow T = T1 ....o T since diffeomorphisms are closed under composition. The flows are trained by maximizing the ligelikli of the observed data under the density (x). The log- lieelikeihood can be evaluated efficiently when it is tractable to compute the Jacobian determinant of T; for example when Ta is a lower triangular function [Marzouk et al., 2016]. Given the map, we can easily generate i..dd. samples from the learned distribution by sampling from the base distribution z' ~ Pz and evaluating the flow T (z').
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 jth component of the flow T has the form: I. = Ti (25; 65 (xxxx())) where each T, an invertible transformer and each Cy is conditioner that only depends on the variables that come before Iy in the ordering т. As a result, the map components define an autoregressive model that factors the density over random variable as: p(x) = I p(=](x²;)) where xzj (x1,.... Fi- autoregressive. Under mild conditions, any arbitrary distribution Px 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 Autoregressive Flows:
Density Estimation with Continuous Normalizing Flows:
Continuous normalizing flows (CNFs) [Chen al., 2018, Grathwohl et al., 2018] represent the transformation T as the flow map solving dz(()
the differential equation Ø8 = [z((t) t; 0). Given the initial condition x = 2(4)) ~ Pxi we can integrate/ / backwards in time from i to lo to obtain
z(to) Px, 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 a Blog p(z()) Øt
Tr( Bu(t) ). This expression used compute the log-likelihood of target sample as logp(z(t1)) = log p(z(t))) - f Tr( (itt)) 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.
8f ddt. We refer the reader to Chen et al. [2018] for the
Causal Inference with Autoregressive Flows: Modelling causal relationships is crucial for en- abling 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 (#....., x2) E iR with joint distribution Pxx the associated SEM consists of d structural equations of the form x, = f, (pay: ut)) where u represents mutually independent latent variables and pa, denotes the direct causal parents of variable z,, Each SEM also corresponds to a directed acyclic graph (DAG), with a causal ordering π defined by the variables' dependencies.
Khemakhem al. [2021] showeased the intrinsic connection between SEMs and autoregressive flows. The authors demonstrated that affine autoregressive flows with 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 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

---

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 approximate complex functions with certain invariances. For a neural network that approximates an arbitrary function f R - R", one type of invariance could consider is when specific output & is independent from given input Iy i.e.:
02g- it 1. x2 means 8e; = 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 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 = (#1,,..,, Tz), we use the lower-triangular adjacency matrix A E {0, 1 to represent the underlying variable dependence structure. In other words, A,j 0 for j < i if and only if æ4 J. 251241.....)1 and Aij = otherwise. This matrix encodes the same information as Bayesian network DAG of the variables. In the fully autoregressive case, matrix A is 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. 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, discussed in Proposition 1.
We improve upon the idea of masked autoregressive neural networks to directly encode the inde- pendence structure represented by 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 autoregressive neural network by factoring the adjacency matrix into binary mask matrices for each hidden layer.
More specifically, given an adjacency matrix A E {0, 1),1) and neural network with L hidden layers, each with h1, ln,,.., hr, hidden units (2 d), we seek mask matrices M E {0, 1 M² e {0, 11) ,...,,M, M E {0, 1jaxhy such that A' ~ A, where A' := M .M.............................. M We use A' ~ A denote that matrices A' and A share the same sparsity pattern, i.e.: exact same locations of zeros and non-zeros. Note that here is a binary matrix and A' is an integer matrix. We then mask the neural network's 's hidden layers using M', .M² ..... M¹ per (1) to obtain Structured Neural Network (StrNN), which respects the prescribed independence constraints. The value of each entry Al thus corresponds to the number connections flowing from input xy output 2; in the StrNN.
Finding the optimal solution to this problem 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 dcconstructing 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 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 lay-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 's approximation error during the matrix factorization step, such 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 ISrivastava 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 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,
d
max A'aA
(2)
max A'a
AlJ var(A"), (3)
j=
where var(AA) 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 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 L optimization involves recursively factorizing the mask matrix layer by layer. Given A E {0. 1]dxa, we run Algorithm I once to find A1 E {0, 1]dxhı and M' {0, 1 xd for layer of width h such that A M ~ A, where ~ denotes that the matrices share the same sparsity. For the next layer with h2 hidden units, we A₁ in place of A to find A2 E (0, 1 1 and M E {0, heah such that A2 M ~ A1. 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. then maximizes the number of neural network connections by filling the second mask with many ones as possible while respecting the sparsity in A. See Appendix A.2 for 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 (dh) time. Scaling up to L layers, where each hidden layer commonly contains (dd) units, the overall runtime is Odd L), which is much more efficient than the integer programming solutions. From our experiments, the greedy algorithm executes nearly instantancously for dimensions in the thousands.
In Appendix A, we include detailed results from in- vestigating the link between the neural network's gen- eralization performance and the choice of mask fac- torization algorithm. We observe that while the exact
Algorithm 1: Greedy factorization
Data: A E {0, 1]dixd2, hidden size h
Result: M¹ E {0, 13d4
solution to objective (2) achieves a higher objective value than the greedy approach, it has no clear advan- tage in density estimation performance. Moreover, we found that models trained with the two objectives, (2)
M E (0, 1 hxd22 satistying
MYM A
no + non-zero rows in A
Fill M with nz: repeat until full
and (3), provide similar performance. However, some
3 Fill M with ones
datasets might be more sensitive to the exact objec-
4 for i-th row in MM do:
tive. For example, problems with anisotropic non- Gaussian 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.
s C' indices of O's in i-th row ofA
T 4- cols. of M whose index E C of T for in R: set M.Y- to zero 9 return (M". .M)
"
R indices of non-zero
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-ting time flows. In the autoregressive flows setting, recall that the jth component of each flow layer is parameterized as zr = (=5; ey (x (X=)())), )), where the conditioner cj 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
5

---

variable is only conditionally dependent on the subset of variables defined by prescribed adjaceney matrix. A single step from StrAF can be represented as: z StrAF(x: (x) T(x, StrNN A)) where x I2 and A {0, (0.1)ax2 1] 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 nor- malizing flows (CNFs), where the neural network that parameterizes f in the differential equation (((()))) is replaced by a StrNN to ob- Bm(e)
Structured Autoregressive Flow (StrAF)
NF Step
Z1
NF Stop
Zz
g
UMNN
UNINN
(x((t) tain 02 StrNN(z((), 1, A). This plug-in replace- ment allows us inject structure into the function describing the continuous dynamics of the CNF with- modifying other aspects of the CNF. We refer this architecture as the Structured CNF (StrCNF).
t 0z
f 0²
StrNN
The StrNNF uses the trace estimator described in FFJORD [Grathwohl et al., 2018] to evaluate the ob- jective function for learning the flow.
True
Connectivity after step
Connectivity
Final Connectivity
Adjacency
after steps
Figure 2: The StrAF injects a prescribed ad- jacency into each flow step using a StrNN conditioner. The StrAF does not permute la- tent variables, allowing the adjacency matrix to be respected throughout the entire flow.
While injecting structure, both StrAF and SuCCNF inherit the efficiency of StrNN due to our choice of weight masking. Specifically, the output of the StrNN can be computed with single forward pass through the network. In comparison, input masking approaches such as Wehenkel and Louppe [2021] must perform forward passes to compute the output for 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 full adjacency matrix. Building on CAREFL, we also assume that the flow T takes on the following affine functional forms for observed data x:
2ラ = (X 2j + tx((xx())) = d
where the functions 8y and Ij are both conditioners that control the dependencies on the variables preceding Ty. It is crucial for the autoregressive or graphical structures to be maintained across all sub-transforratiins 1,,.. ..., Tk for the flow T2
Tk.
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 additional mask- ing 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 masking scheme similar to methods in Germain et al. [2015]. We again note that input masking scales poorly with data dimension d, forward passes are required to obtain the output for a single datum.
Silvestri al. [2021] proposed embedded- model flows, which alternates between traditional nor- malizing 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 cor- rect. 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 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, Weilhach 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 apply 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 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
Flows have garnered increasing interest in the context of causal inference, with applications spanning various problem domains. Ilse 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 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 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 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 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://// //gi thub. cwmm http:/gkkabbbburgg /Sttruuuturrdll.e..
density
5.1 Density Estimation 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-dinml 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 square of a pixels determined by the hyper-parameter By changing the hyperparameter we can increase the context window used to model each pixel.
7

---

Synthetic: Binary every_other other
Synthetic: Binary random, sparse
Binary MNIST Label
Binary Label
29
250
MADE StrMN
MADE StNN
150
MADE StrNN
1550
MADE
StrNN
225
28
140
200
27
130
130
120
120
26
150
110
110
2000 4000 6000 8000 10000 Semple Size
0
2000
4000
500
1000
1500
2000
500
1000
1500
2000
Sample Size
Sample Size
Semple Size
Figure 3: Negative log likelihood of test set (lower is better) for density estimation experiments binary synthetic data (left 2 images) and label-dependent binary MNIST data (right 2 images). Error ranges reported as standard error across the test set. StrNN performs better than MADE as the sample size decreases.
StrNN is equivalent to MADE for this experiment when we set nbr _or__seeee.. 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 Gaussian Data
Synthetic: Gaussian
Synthetic: Gaussion random
We run experiments to compare the perfor- mances of StrNN and MADE on synthetic Gaus- sian data generated from known structure equa- tion models where each 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 density estimation than the fully autoregressive MADE baseline.
MADE SUINN
4.5
MADE
-13
-15
3.5
-16
3.0
2.5
2000
4000
1000
2000
3000
4000
Samble Size
Samole Suz
Figure 4: Results from density estimation experiments on Gaussian synthetic data generated from two diffcrent sparsity patterns. Test loss is reported in negative log likelihood with error ranges (standard error across test set). StrNN performs significantly better than MADE in the low data regime, and better average.
5.3 Density Estimation with Structured Normalizing Flows
We evaluate StrAF on density estimation against several baselines. We draw 1000 samples from 15 dimensional tri-modal and non- lonelr synthetic dataset for experimental evaluation. Data generation 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] 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-hased models for density estimation
We report results in Table and observe several trends. For both discrete and continuous flows, the abil- ity 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 struc- ture. For example, while the method proposed by Weilbach et al. [2020] allows DAG structure to
8

---

ARF
GNF
FFJORD
Weilbach
StrAF
StrCNF
(g)
X12)
(Xa. X20)
Figure 5: Model generated samples shown in blue dots for randomly selected dimensions. The ground truth density visualized by the orange contours. See Appendix E.3 an explanation on why GNF performs poorly.
be injected into the baseline CNF, its inability to easily use multi- layered 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.
Table 1:
Evaluation of flow- based models. Mean and 95% CI of test NLL over 8 runs reported.
Test NLL (4)
ARF GNF
-3.09 0.43
-3.63 + 0.35
StrAF
-3.55 + 0.20
5.4 Causal Inference with Structured Autoregressive Flows
FFJORD
-1.85 + 0.64
Weilbach
-2.59 + 0.58
We conduct synthetic experiments where 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 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 Appendix D.3. addition, Appendix D. provides detailed algorithms and related discussions on generating interventional samples and computing counterfactuals with flows.
SuCNF
-4.01 + 0.12
S-variable SEM
10 ID-variable SEM
S-variable SEM
10 10-variable SEM
1.5
2.5
CAREFI StrAF
CARER.
15
CARER. StoAF
10tal I-NSE
3.0
StrAF
1.0
Total CIMSE
StrAF
0.5
0.0
0.0
500
1000
1500
2000
500
1000
1500
2000
500
1000
1500
2000
5000
1000
1500
2000
Sample size
Sample size
Sample size
Sample size
(a) Interventions
(b) Counterfactuals
Figure 6: Evaluations of causal predictions (left: interventions: right: counterfactuals) on 5- and 10- variable SEMs made by StrAF and CAREFL. Performance is measured by the corresponding total mean squared error with standard deviation across ten runs. (a) measures the error of the expected value of one variable under different interventions, while (b) computes the error by deriving counterfactual values under different observed samples queries.
Interventions. perform interventions on each variable 2; 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 E[z,/do(z)) z a)] for variable 24, excluding the intervened variable z, 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. Coun- terfactual inference tackles what-if scenarios: determining the value of variable z; if variable zy had taken a different value a. Unlike interventions, counterfactual queries involve deriving the latent variables z given the observed data x = Xobs, rather than sampling new z. We generate 1000 observations using the synthetic SEM and derive counterfactual values x by posing queries with varying a values for each variable zj. 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 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 fully autoregressive MADE model on synthetic and MNIST data. We integrate SurNN 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. We assume access 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 Icarned 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 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 valuable direction [Frankle and Carbin, 2018). 8]. 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-Mereshh 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 CGPIN-2022-04546 and a Canada CIFAR AI Chair. AC is supported by 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 part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.
10

---

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 PMLR, 2021.
modeling. In International Conference Machine Learning, pages
Limeng Cui, Haeseung Seo, Maryam Tabar, Fenglong Ma, Suhang Wang, and Dongwon Lec. 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. 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.081.3, 2022.
Esteban G Tabak and Cristina V Turner. A family of nonparametrri 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.
Mathicu Germain, Karol Gregor, lain Murray, and Hugo Larochelle. MADE: Masked autocncoder for distribution estimation. In International Conference on Machine Learning, pages 881-889. PMLR, 2015.
Judea Pearl. Bayesian networks. 2011.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation. Advances neural information processing systems, 30, 2017.
Chin-Wod 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): 2021.
Youssef Marzouk, Tarek Moselhy, Matthew Parn,, and Alessio Spantini. An introduction to sampling via measure transport. arXiv preprint arXhv: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. Gamett, ed- itors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL https: https://proceediing.. neurips. wwtppppepeaiiiieeeceee semper_ffless files/ille/e paaape__iieee/ppprr/20000009 pttpppppee/////leee 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 dynamies for scalable reversible generative models. arXiv preprint arXiv:1810.01367. 2018.
11

---

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 uXXi::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 Mathe- maticae (Proceedings), 90(5):406-422) 1977. ISSN 1385-7258. doi: https://doi.org// 1016/ 1388-7255877))9005555 URL httpps//wwat sciencedirect. wwt.://eeeerrtecceeeeeeeeeeeeeee wwt.sccic/cee//eeecceccccoeccece htt.n///eeriieeeppiii.eeeee 1385725877900555..
Chen Dan, Kristoffer Arnsfelt Hansen, He Jiang, Liwei Wang, and Yuchen Zhou. Low rank ap- proximation of binary matrices: Column subset selection and generalizations. arXiv preprint arXiv:1511.0169, 2015.
Fedor 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: 880300363, 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()):1929-1988, 2014.
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncer- tainty 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 mlr. press/v485 gal16. html.
Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL https://www.
gurobi. com. 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 preprini arXiv:2110.00
2021.
Jacobie Mouton and Steve Kroon. Graphical residual flows. arXiv preprint arXiv:2204. /1846 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: httht::s//prcceeiing mlr. press/v108/ weilbach20a. html.
Francois Rozet, Felix Divo, and Simon Schnake. Zuko: Normalizing Flows in PyTorch. htwp://thwccoooemmmoooooooreeeee
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.06200 2022.
12

---

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,
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-MICCAl 2021: 24th International Conference, Strasbourg, France, Seplember 27-October 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 cousal-graphicll normalizing flows. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 18 11880-111 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.0939), 2022b.
Ruslan Salakhutdinov and lain 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 Xing. Dags with no tears: Continuous optimization for structure learning. Advances in neural information processing systems, 31, 201 8.
Diviyan Kalainathan, Olivier Goudet, and Ritik Dutta. Causal discovery toolbox: Uncovering causal relationships in python. The Journal of Machine Learning Research, 21(1):1406-14D) 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.
13


Reviews:
[{"review_id": "PbYfBAZFXm", "reviewer": "Reviewer_Tj4q", "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. \nWith 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. ", "questions": "1) How GNF would compare to StrAF if it does permute the latent variables after the first step? \n2) 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. \n", "limitations": "N/A", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n2- The paper didn't propose any approach for learning the structure of DAG given the provided structured neural network parametrization.\n3- MADE is not a strong density estimator and comparing only to MADE does not validate the strength of structured neural networks as density estimators. \n ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer Tj4q for their feedback and insightful comments, especially for highlighting the simplicity of our matrix factorization approach and our contribution to the causal inference application.\n\nFirst, we hope to provide some further clarification on the questions posed by the reviewer:\n\n> How GNF would compare to StrAF if it does permute the latent variables after the first step?\n> \n\n**We believe GNF performs worse than StrAF when the latent variables are permuted after the first step.** The GNF already permutes latent variables after the first step, but does not address the effect of this operation on the conditional independencies enforced by the overall network (i.e., the resulting adjacency matrix between inputs and outputs). This is consistent with how the GNF authors have implemented their official code-base (see https://tinyurl.com/vs92kncc). As we visualize in Figure 2, this permutation can encode incorrect connectivity between the variables. Our experiments show this causes higher MMD values between samples from GNF and the true distribution, in comparison to StrAF.\n\n> During the comparison with CAREFL did you provide CAREFL with external DAG orders that StrAF uses? …\n> \n\n**Yes, we provide the ground truth DAGs to both StrAF and CAREFL during the evaluations for causal inference.** However, CAREFL can only utilize the autoregressive order provided whereas StrAF enforces the entire adjacency structure, which results in improved performance. Even though CAREFL discussed approaches for extending their pairwise discovery method to multivariate data, they assumed the autoregressive orders are given in their causal inference experiments. Similarly, one major assumption in our experiments is that the true structure has been resolved before modelling the generative process.\n\nNext, we address some of the weaknesses of our paper as identified by the reviewer.\n\n> 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.\n> \n\nWe believe that, **while StrNN is similar to MADE, the augmentation itself makes a significant  contribution to the literature** for the following reasons:\n\n1. While MADE identifies an approach to enforce autoregressive dependence in NN outputs by masking weights, their algorithm does not consider sparse autoregressive dependence, which appears in many applications such as causal inference. \n2. The way we enforce structure through matrix factorization allows us to directly specify a neural architecture. As far as we know, this important link between masking and NN architecture has not been explicitly discussed or explored in MADE or other prior works. We refer reviewer Tj4q to our overall author response for additional experiments on how the factorization objectives can impact NN performance.\n\n> The paper didn't propose any approach for learning the structure of DAG given the provided structured neural network parametrization.\n> \n\n**In the conclusions and limitations section of our manuscript, we mention that while we assume access to the true adjacency matrix, one can use any SOTA structure discovery algorithm, such as NO TEARS, in conjunction with StrNN and StrAF, as long as one keeps the flow layers un-permuted.** We will also add an additional reference to the causal-learn [1] package, which provides a variety of out-of-the-box structure discovery algorithms that can be easily integrated into our StrNN/StrAF code.\n\n**While structure discovery is outside the scope of this paper, we are interested in pursuing future work in using StrNN to learn structure from data, providing a full pipeline from structure discovery to density estimation and sample generation.** One option is to add a group regularization penalty to the network masks or weights in order to penalize the dependence of certain marginal conditionals on individual input variables. We have also thought of an alternative stochastic approach to structure discovery, by using dropout to reduce the total number of connections, which is similar to the proposal from reviewer Wq9i based on the lottery ticket hypothesis.\n\n> MADE is not a strong density estimator and comparing only to MADE does not validate the strength of structured neural networks as density estimators.\n> \n\nWe compare to MADE as a baseline because it is the natural choice that assumes a full autoregressive factorization. We do not claim that StrNN is the strongest binary density estimator, rather we show that **adding structure to several existing methods improves their performance.**\n\nWe later show that incorporating StrNN into popular continuous-valued density estimation methods based on autoregressive flows improves their performance. Our baseline, GNF, is a very strong density estimator, and Table 1 in our paper shows that StrAF can improve upon GNF for both test likelihood values and sample quality.\n\nTo reinforce this point, we incorporated StrNN with continuous normalizing flows, FFJORD. We report results in the overall author response. We observe that adding StrNN improved performance against the baseline FFJORD model, again demonstrating the strength of StrNN, not as an independent density estimator, but as a drop-in replacement to improve the performance of density estimators by leveraging adjacency structure. We aim to further study StrNN to improve current SOTA density estimators, such as diffusion models.\n\nLastly, we thank Reviewer Tj4q once again for your constructive feedback. We are confident our revised manuscript can address your questions and concerns, and we kindly invite you to reconsider the score for our paper if you are satisfied with our responses.\n\n[1] Zheng Y, Huang B, Chen W, Ramsey J, Gong M, Cai R, Shimizu S, Spirtes P, Zhang K. Causal-learn: Causal Discovery in Python. arXiv preprint arXiv:2307.16405. 2023 Jul 31."}, {"review_id": "bIEUuKWMtk", "reviewer": "Reviewer_Wq9i", "summary": "In 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. ", "questions": "As 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", "limitations": "The 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.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "I 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.", "weaknesses": "The 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to warmly thank Reviewer Wq9i for their encouraging feedback, especially for highlighting our method’s emphasis on data efficiency and the novelty of the causal effect estimation application.\n\nThe reviewer raises a good point on how broader comparison on other datasets would make our paper more robust. While the synthetic datasets included in our evaluation are highly non-linear and challenging to model, we are working on extending our method to several real world datasets. We have obtained favorable density estimation results in adding StrNN to continuous normalizing flows such as FFJORD, which we have included in the general author response. This further validates StrNN’s efficacy as a density estimation technique.\n\n**re: StrNNs and structure discovery**\n\n> Explore the possibility of discovering the dependencies for a given order. Here 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'm wondering if you explored similar ideas during your research?\n> \n\nThat is a great point! Overall, structure discovery is outside the scope of our paper since we assume either the ground truth Bayesian network or a good estimate is known. However, **we are interested in pursuing future work in using StrNN to learn structure from data, thus providing a full pipeline from structure discovery to density estimation and sample generation.**\n\n**We have considered the following potential connections to dropout and the lottery ticket hypothesis during our work.** \n\nFirst, we noticed that stochastically introducing sparsity into a neural network has interesting connections to dropout, which is also known to improve model generalization. It would be interesting to investigate if specific patterns of sparsity emerge from this stochastic process, and as the reviewer points out, if our method provides a framework with which to directly optimize for any specific qualities of sparsity while respecting structure.\n\nSecond, we have considered a gradient-based approach to learn structure from data as well, where we might add a group regularization penalty to the weights in order to penalize the dependence of certain marginal conditionals on individual input variables.\n\n> The 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> \n\nLastly, we thank the reviewer for raising an important point on the core limitation of all autoregressive models, which is the difficulty of performing marginal and MAP inference. We will make this point clearer in our final revision.\n\nWe would like to thank Reviewer Wq9i once again for your constructive feedback. We are confident our revised manuscript can address your questions and concerns, and we kindly invite you to reconsider the score for our paper if you are satisfied with our responses. Please don’t hesitate to let us know of any additional comments!"}, {"review_id": "JZNZGhdhVG", "reviewer": "Reviewer_icwN", "summary": "The 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.", "questions": "1. 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.", "limitations": "No concerns here", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer icwN kindly for their encouraging and constructive feedback, especially for highlighting the novelty and efficacy of our matrix factorization approach to enforcing exact independence structures in arbitrary NNs.\n\n> The example shown in Figure 1 decomposes the network into two separate networks.\n> \n\nWe will use an alternative example adjacency matrix in Figure 1 which requires shared weights in our final revision. We thank the reviewer for their attention to detail on this matter.\n\nA main point in the reviewer’s feedback is regarding our design choice to use one single neural network with shared weights in StrNN. In order to enforce independence structure among variables in a single NN, we are required to solve the binary matrix factorization problem central to our paper. The reviewer is correct in saying we can instead find a separate NN for each output variable, which only takes that variable’s dependencies as inputs. That approach would be very simple and render masking and matrix factorization procedures unnecessary. However, W**e believe that the particular way we share weights in StrNN has several advantages**:\n\n1. **Sharing weights is more efficient for improving both memory and computational efficiency against other approaches for masking.** For example, when there are D input features, using D independent networks or the GNF baseline (which masks input variables) necessarily requires D forward passes for a single input example. In comparison, we can do the same computation in a single forward pass. This becomes especially relevant in contexts such as continuous normalizing flows (CNFs), where many network evaluations must be done in the ODE solver. Here, methods with poor asymptotic complexity become impractical. In our additional results for integrating StrNNs into CNFs with the dataset from Section 5.3 (as described in the **general author response**), we find that the StrNN model outperforms a baseline FFJORD even when using a number of hidden units on the same order of magnitude (as opposed to D times more,  see **supplementary rebuttal PDF** for exact experiment details). This would not be possible unless the model leveraged weight sharing in StrNNs,  providing further empirical justification for the utility of this choice.\n2. **From a causality perspective, one possible benefit from sharing weights could come from robustness to model misspecification. T**he best model depends on the actual problem at hand, its associated data, and possibly the true generative process. To our knowledge, an empirical study comparing sharing weights vs independent models does not exist and would be interesting for the community. Further investigation into this question using the StrAF extension of CAREFL we built in Section 5.4 of our paper would be a promising direction for future research. However, as suggested in [1], if the functional forms of the independent mechanisms are misspecified, then variables outside of the Markov blanket may improve the predictive model performance. By sharing weights inside the network, gradient information from other variables’ relationship can improve generalization of a given conditional distribution. \n\nLastly, the reviewer asked us to discuss\n\n> possible simplification of specific causal structures, i.e., independent causal mechanisms as in Figure 1.\n> \n\n**We believe this is a fruitful research direction of applying our strNN framework to other statistical models in causal inference that may inherit sparse structure.** Two possible simplifications are:\n\n1. **Additive noise models with a common noise mechanism**, i.e., the data-generating process for each variable $x_k$ is given by $x_k = f_k(x_1,\\dots,x_{k-1}) + eps$, where eps is independent of k. In this case, the strNN outputs will only need predict the conditional expectation $f_k = \\mathbb{E}[x_k | x_1,\\dots,x_{k-1}]$ for all k, and a separate network is used to model the distribution of the noise variable $eps$. \n2. **Graphical models with sub-groups of connected components**. For example, Figure 1 shows a model with four variables where (x_1,x_2) is a connected group, and (x_3,x_4) is connected to the first group by an edge from x_2. In this case, we could apply strNN to groups of variables in a hierarchical way: first to the sub-group using a single edge between then, and then to the variables within each group. \n\nFinally, we would like to extend our thanks once again to Reviewer icwN for the insightful feedback provided. We believe that the changes made to our manuscript can address your questions and concerns. We hope you will find the revisions satisfactory and consider a reevaluation of our paper's score. Should you have any further feedback or questions, please feel free to share.\n\n[1] Peters J, Janzing D, Schölkopf B. Elements of causal inference: foundations and learning algorithms. The MIT Press; 2017. Chapter 6. Multivariate Causal Models. p.103."}, {"review_id": "rCmtiMwN4w", "reviewer": "Reviewer_LoZH", "summary": "This 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.", "questions": "* 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?", "limitations": "* 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.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* 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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer LoZH for their detailed feedback, especially for highlighting the clarity and soundness of our approach to enforce conditional independences in NNs.\n\nWe begin with the Zuko repo, which is relevant to our work. First, are there manuscripts or research results using Zuko that we could cite? We could not find any based on the GitHub authors' profiles. Otherwise we will reference the repo in our revision.\n\nSecond, we note that the **code in the Zuko repo (as cloned on the NeurIPS submissions deadline)** **does not allow users to initialize any of its flows to encode a global adjacency structure**, and attempting to do so gives an error. \n\nThird, **while Zuko and StrNN’s greedy factorizations of the adjacency matrix are similar, they yield different results.** In the experiment below, mask matrices found by our greedy algorithm and by Zuko result in different performance. Zuko operates on the unique rows of an adjacency matrix, which causes issues when the matrix contains many repeating rows. This type of structure may be naturally encountered in datasets with star shaped graphs. Consider a NxN matrix where the first (N-1) rows contain dependence on the first variable ([1, 0, …, 0]) and only the last row depends on the second variable ([0, 1, 0, …, 0]). Given a budget of H hidden units, Zuko would assign H // 2 units to represent the variable corresponding to the last row. This ineffectively represents the dependence of all other outputs on the first variable, and is avoided by our greedy algorithm. In a non-linear dataset with N variables and the above adjacency matrix, we compare StrAF with our greedy algorithm and Zuko’s algorithm in the table below. We report runs from 5 random seeds, using a 95% CI and N hidden units.\n\n| Model | Test NLL (N=50) | Test NLL (N=1024) |\n| --- | --- | --- |\n| StrNN-Greedy | 0.4311 ± 0.5406 | -18.734 ± 1.999 |\n| Zuko | 1.3276 ± 0.2806 | -14.550 ± 0.936 |\n\nWe see our method performs better, especially as N increases. **This supports our claim that the choice of mask factorization can impact performance beyond enforcing independences.**\n\nFourth, Wehenkel and Louppe (2021) and Lachapelle et al (2019) did suggest a masking scheme like MADE. **But there is a gap between the suggestion of an idea and a practical instantiation thereof.** Neither works studied any concrete algorithms or experiments to this end, and only raised the possibility of such an approach for future work. **We believe our work establishes the feasibility of the method, conducts rigorous evaluations, and demonstrates its application in a novel context.** \n\nWe reply to other questions and concerns below:\n\n> Algorithm 1: … how the … algorithm is applied when ... StrNN has more than one hidden layer.\n> \n\n**We apply the factorization algorithm recursively, layer-by-layer, for efficiency and simplicity.** Given $A\\in \\{0, 1\\}^{d \\times d}$, we run Algorithm 1 once to find $ A_1 \\in \\{0, 1\\}^{d \\times h_1}$ and $ M^1 \\in \\{0, 1\\}^{h_1 \\times d}$ for a layer of width $h_1$ such that $A_1\\cdot M^1 \\sim A$, where $\\sim$ denotes the matrices share the same sparsity. For the next layer with $h_2$ units, we use $A_1$ in place of $A$ to find $A_2 \\in \\{0, 1\\}^{d \\times h_2}$ and $M^2 \\in \\{0, 1\\}^{h_2 \\times h_1}$ such that $A_2 \\cdot M^2 \\sim A_1$. Repeat until we have all the masks.\n\n> Section 5.1: same number of layers/neurons used for StrNN and MADE …?\n> \n\n**Yes**, the NNs in StrNN and MADE have the same number of layers/hidden units to control for capacity.\n\n> Line 317: \"As GNF permutes variables …\". … not able to find a mention of this in [25].\n> \n\n**See the following link (with line #) in the GNF repo where the authors permute variables between flow steps:** https://tinyurl.com/vs92kncc. Variable permutation is common in the normalizing flow literature. We believe pointing out the necessity of properly handling adjacency masks with respect to this operation to be a simple yet impactful observation.\n\n> Figure 5/Table 1:  how \"ARF-10\" and \"GNF-10\" can be so much worse than \"GNF-1\" …\n> \n\n**High expressivity of a model does not necessarily lead to better generalization.** Though we have conducted hyper-parameter/model selection for all methods, as detailed in Appendix E.2, the baseline models can still over-fit on spurious correlations. When an adjacency matrix is known, we can avoid this, leading to the improvement of GNF-1 over ARF-10. As shown in Figure 2 and the link above, GNF does not properly handle variable permutations and adjacency masking, leading to an incorrect global adjacency matrix, hence the drop in performance between GNF-1 and GNF-10. We will edit our final revision for clarity on this matter. \n\n> Table 1: … distinction between \"density estimation\" and \"sample quality\"…\n> \n\nWe thank the reviewer for allowing us to clarify this subtle point. This distinction arises from the metric used to assess how *well* the flow estimates the target density. Common divergences between probability distributions used to learn the flow, e.g. KL divergence, are weighted by the target distribution. Thus, the loss is not very sensitive to the model’s performance in low-probability regions of the target. This behaviour might lead to poor “sample quality” even if the likelihood values on the test set is high. \n\n> … StrNN in other settings than density evaluation, such as physics-informed machine learning …\n> \n\nThis is a great direction for future research! In our paper, we studied causal effect estimation as a novel application area. **We will keep physics-informed ML in mind for future work.**\n\n> (StrNN) does not apply to convolutional, attention-based or recurrent networks, .. does not support skip/residual connections or normalization layers.\n> \n\n**StrNN can be extended to several of these architectures.** While our focus is on MLPs, the StrNN can be applied in the linear layers of recurrent networks to enforce structure throughout the sequence. **This is beyond the scope of this work.**"}]

Author rebuttal:
We thank all reviewers for your engagement and thoughtful comments. We have addressed specific concerns and questions in individual rebuttals, but here we **highlight two experiments added based on the feedback**:

1. We extend Section 5.3 with experiments on **continuous normalizing flows (CNF)**. We choose FFJORD [1] (denote as **FFJORD-FC**), a popular strong density estimator, as our baseline, and improve it by injecting prior known structure via StrNN. In particular, we replace the feed-forward network used to model the flow dynamics with a StrNN, denoted **FFJORD-StrNN**. We also compare against [2], which also attempts to encode structure into the ODEnets of a FFJORD model (denote as **FFJORD-Weilbach**). However, FFJORD-Weilbach does not use mask factorization and directly multiplies an adjacency matrix onto a single hidden layer. Thus, they are restricted to a single hidden layer with same number of units as the input dimension, causing underperformance. We note that [2] proposes interesting orthogonal avenues of improving CNFs which we can explore in the future. Using the dataset from Section 5.3, we perform a hyperparameter search (**Table 1** in PDF document). We report Test NLL with a 95% CI based on 8 runs from random initializations below:
    
    
    | Model | Test NLL |
    | --- | --- |
    | FFJORD-FC | -3.9718 ± 0.0677 |
    | FFJORD-StrNN | **-4.0583 ± 0.0089** |
    | FFJORD-Weilbach | -1.0624 ± 0.0724 |
    
    Integration of StrNN results in significantly better performance against both baselines. **We believe this reinforces the comments identifying StrNN as a generally applicable, drop-in tool to improve existing density estimators.**
    
2. We highlight the importance of the StrNN capability to **specify different objectives in the adjacency matrix factorization algorithm**. In particular, we wish to demonstrate the consequences of sub-optimal factorizations. We show density estimation performance using two alternative factorizations: (1) **mincon**, which minimizes the number of connections while remaining faithful to the adjacency matrix, and (2) **random**, which recreates the random factorization scheme from [3]. We compare these two against our **greedy** algorithm when initializing StrNNs in a StrAF model. Our dataset is based on a 10-dimensional sparse linear additive SEM, similar to Section 5.4. We use a 5-step flow with 3 hidden layers each with 20 hidden units. We report test NLL with a 95% CI based on 5 runs from different random initializations. We also report the fraction of remaining connections in the masked NN averaged across layers, where 1.0 is equivalent to a fully connected network, and the test NLL when the flow is trained using N samples.
    
    
    | Factorization Method | Fraction remaining connections | Test NLL (N=100) | Test NLL (N=800) |
    | --- | --- | --- | --- |
    | Greedy | 0.255 | **15.765 ± 0.066** | **14.960 ± 0.115** |
    | Random [3] | 0.1245 | 16.733 ± 0.291 | 15.376 ± 0.183 |
    | Mincon | 0.085 | 17.607 ± 0.222 | 17.094 ± 0.078 |
    
    We also visualize the trend of these results as a function of dataset size (Figure 1 in the attached PDF). **We see clear improvement using the greedy approach, reiterating the importance of factorization with respect to an objective.** 
    

Furthermore, we clarify **common concerns** the reviewers have raised:

1. **Comparisons to MADE**: Reviewers noted that we drew inspiration from MADE to mask weight matrices. We believe **our extension of this idea to enforce a specific, sparse adjacency structure is novel**. Further, a significant contribution is our observation that the **objective of the mask factorization is directly linked to NN architecture, which can affect performance**, as demonstrated in our experiments above.
2. **Alternative methods for enforcing structure (such as GNF):** Reviewers noted that other methods of enforcing a prescribed adjacency structure exist. Given D-dimensional input, one could use GNF, or initialize D independent neural networks, each accepting the dependent features for a specific variable. However, both these approaches require D forward passes to compute the output of a single input vector, whereas we only require one. **The lower asymptotic complexity of evaluating our model enables higher training efficiency** and is important when incorporating StrNN into density estimators such as CNFs (see above), which require many forward evaluations to compute a single ODE solution.
3. Several reviewers correctly noted that our current work on StrNN does not have the capability to **learn Bayesian networks from data**. **This was by design given the main assumption in our work was that we would have access to the ground truth DAG, whether from domain knowledge or prior experiments with structure discovery algorithms.** To direct readers to these resources, we will add citation to the **causal-learn package** in our final revision. We have further considered different ways we could utilize StrNN itself to learn structure from data, as discussed in our response to Reviewer Tj4q.

Finally, we would like to thank all reviewers once again for your insightful responses, which have helped us refine our work and identify future directions. We are confident our updated manuscript will have addressed your concerns and suggestions, and we kindly invite you to reconsider our paper if you are satisfied with our responses.

[1] Grathwohl, Will, et al. Ffjord: Free-form continuous dynamics for scalable reversible generative models. *International Conference on Learning Representations*, 2019.

[2] Weilbach, Christian, et al. "Structured conditional continuous normalizing flows for efficient amortized inference in graphical models." *International Conference on Artificial Intelligence and Statistics*. PMLR, 2020.

[3] Mouton, Jacobie, and Steve Kroon. "Graphical residual flows." *arXiv preprint arXiv:2204.11846* (2022).

Meta-review: A good contribution to combine neural networks and graphical models such as Bayesian networks and causal models. The reviewers generally favored acceptance. Some concerns by the reviewers remain, which should be addressed by the authors.

The AC has some additional comments:

Since binary matrix factorization is NP-hard, this means that Algorithm 1 does not necessarily find a solution with the exact sparseness patterns prescribed by $A$, correct? As far as I get it, the Algorithm might introduce further zeros in $A'$, is this correct? The authors should please clarify this in the paper.

Please double check and enhance the description of Algorithm 1.

- First, it was not clear what line "Fill $M^W$ with $nz$; repeat until full" means. I only got it from Figure 7a in the appendix.
- In the visual example, perhaps start with $M^V$ indeed filled with ones, to see the effect of the algorithm.
- In line 4, we iterate $i$ over the rows of $M^V$. Then we look for the indices in the $i$-th row in $A$. However, A has d1 many rows and $M^V$ has d2 many rows, thus we potentially address into non-existing rows of $A$, if $d1 < d2$. (I suppose the algorithm should also work for $d1 \not= d2$?)
Decision: Accept (poster)