SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the general real-valued probabilistic density estimation task, we highlight the use of StrNN in two popular normalizing flow frameworks, autoregressive flows and continuous-time flows. In the autoregressive flows setting, recall that the  $j$ th component of each flow layer is parameterized as  $x_{j} = \tau_{j}(z_{j};c_{j}(\mathbf{x}_{<\pi(j)}))$ , where the conditioner  $c_{j}$  dictates which inputs the latent variable can depend on. We use the StrNN as the conditioner, combined with any valid invertible transformer  $\tau$ , to form the Structured Autoregressive Flow (StrAF). The StrNN conditioner ensures each transformed latent
variable is only conditionally dependent on the subset of variables defined by a prescribed adjacency matrix. A single step from StrAF can be represented as:  $\mathbf{z} = \mathrm{StrAF}(\mathbf{x}) = \tau (\mathbf{x},\mathrm{StrNN}(\mathbf{x},A))$  where  $\mathbf{x}\in \mathbb{R}^d$  and  $A\in \{0,1\}^{d\times d}$  is an adjacency matrix. Normalizing flow steps are typically composed to improve the expressiveness of the overall flow. To ensure that the prescribed adjacency structure is respected throughout all flow steps, we simply avoid the common practice of variable order permutation between flow steps. We visualize this in Figure 2. This choice does not hurt performance as long as a sufficiently expressive transformer is used. In our experiments, we apply the unconstrained monotonic neural network (UMNN) described in Wehenkel and Louppe [2019] as a transformer.
We also integrate the StrNN into continuous normalizing flows (CNFs), where the neural network that parameterizes  $f$  in the differential equation  $\frac{\partial \mathbf{z}(t)}{\partial t} = f(\mathbf{z}(t), t)$  is replaced by a StrNN to obtain  $\frac{\partial \mathbf{z}(t)}{\partial t} = \mathrm{StrNN}(\mathbf{z}(t), t, A)$ . This plug-in replacement allows us to inject structure into the function describing the continuous dynamics of the CNF without modifying other aspects of the CNF. We refer to this architecture as the Structured CNF (StrCNF). The StrCNF uses the trace estimator described in FFJORD [Grathwohl et al., 2018] to evaluate the objective function for learning the flow.
While injecting structure, both StrAF and StrCNF inherit the efficiency of StrNN due to our choice of weight masking. Specifically, the output of the StrNN can be computed with a single forward pass through the network. In comparison, input masking
approaches such as Wehenkel and Louppe [2021] must perform  $d$  forward passes to compute the output for a single datum. This not only prevents efficient application of input masking to high dimensional data, but also is a barrier to integrating the method with certain architectures. For example, the CNF already requires many neural network evaluations to numerically solve the ODE defining the flow map, so making  $d$  passes per evaluation is particularly inefficient.
Structured Autoregressive Flow (StrAF)
Figure 2: The StrAF injects a prescribed adjacency into each flow step using a StrNN conditioner. The StrAF does not permute latent variables, allowing the adjacency matrix to be respected throughout the entire flow.
3.3 Structured Causal Autoregressive Flow
Beyond density estimation, StrAF allows us to build autoregressive flows that faithfully represent variable dependencies defined by a causal DAG. In contrast, CAREFL [Khemakhem et al., 2021] only maintains the autoregressive order, which is insufficient for data with more than two variables where the true structure must be characterized by a full adjacency matrix. Building on CAREFL, we also assume that the flow  $\mathbf{T}$  takes on the following affine functional forms for observed data  $\mathbf{x}$ :
$$
x _ {j} = e ^ {s _ {j} \left(\mathbf {x} <   \pi (j)\right)} z _ {j} + t _ {j} \left(\mathbf {x} <   \pi (j)\right), \quad j = 1, \dots , d \tag {4}
$$
where the functions  $s_j$  and  $t_j$  are both conditioners that control the dependencies on the variables preceding  $x_j$ . It is crucial for the autoregressive or graphical structures to be maintained across all sub-transformations  $\mathbf{T}_1, \ldots, \mathbf{T}_K$  for the flow  $\mathbf{T} = \mathbf{T}_1 \circ \dots \circ \mathbf{T}_K$ .
Assuming the true causal topology has been given either from domain experts or oracle discovery algorithms, StrAF can directly impose the given topological constraints by adding an additional masking step based on the adjacency matrix, as shown in section 3.1. This ensures that the dependencies of the flows match the known causal structure, which leads to more accurate inference predictions.
Given a Bayesian network adjacency matrix, Wehenkel and Louppe [2021] introduced graphical conditioners to the autoregressive flows architecture through input masking. They demonstrated that unifying normalizing flows with Bayesian networks showed promise in injecting domain knowledge while promoting interpretability, as even single-step graphical flows yielded competitive results in density estimation. Our work follows the same idea of introducing prior domain knowledge into
autoregressive flows, but we instead use a masking scheme similar to methods in Germain et al. [2015]. We again note that input masking scales poorly with data dimension  $d$ , as  $d$  forward passes are required to obtain the output for a single datum.
Silvestri et al. [2021] proposed embedded-model flows, which alternates between traditional normalizing flows layers and gated structured layers that a) encode parent nodes based on the graphical model, and b) include a trainable parameter that determines how strongly the current node depends on its parent nodes, which alleviates error when the assumed graphical model is not entirely correct. In comparison, our work encodes conditional independence more directly in the masking step, improving accuracy when the assumption in the probabilistic graphical structure is strong.
Mouton and Kroon [2022] applied a similar idea to residual flows by masking the residual blocks' weight matrices prior to the spectral normalization step according to the assumed Bayesian network. Similarly, Weilbach et al. [2020] introduced graphical structure to continuous normalizing flows by masking the weight matrices in the neural network that is used to parameterize the time derivative of the flow map. However, their method is difficult to apply to neural networks with more than a single hidden layer. In comparison, our factorization algorithm more naturally permits the use of multi-layered neural networks when representing CNF dynamics. The Zuko software package [Rozet et al., 2023] implements various types of normalizing flows, including MAFs that also rely on weight masking, but they only provide one possible algorithm to enforce autoregressive structure given a specific variable order. In comparison, our approach permits explicit optimization of different objectives during the adjacency matrix factorization step and we investigate the efficacy of factorization schemes and resulting neural architectures in our work. For completeness, we include pseudocode of their algorithm and comparisons to our own algorithms in Appendix A.
Flows have garnered increasing interest in the context of causal inference, with applications spanning various problem domains. Ilse et al. [2021] parameterized causal model with normalizing flows in the general continuous setting to learn from combined observational and interventional data. Melnychuk et al. [2022] used flows as a parametric method for estimating the density of potential outcomes from observational data. Flows have also been employed in causal discovery [Brouillard et al., 2020, Khemakhem et al., 2021] as well as in various causal applications [Ding et al., 2023, Wang et al., 2021]. In particular, Balgi et al. [2022a] also considered embedding the true causal DAG in flows for interventional and counterfactual inference, but they do so via the framework of Graphical Normalizing Flows [Wehenkel and Louppe, 2021]. Balgi et al. [2022b] used CAREFL on a real-world social sciences dataset leveraging a theorized Bayesian network.
To demonstrate the efficacy of encoding structure into the learning process, we show that using StrNN its flow integrations to enforce a prescribed adjacency structure improves performance on density estimation and sample generation tasks. We experiment on both synthetic data generated from known structure equations and MNIST image data. Details on the data generation process for all synthetic experiments can be found in Appendix C. We also apply StrAF in the context of causal inference and demonstrate that the additional graphical structure introduced by StrAF leads to more accurate interventional and counterfactual predictions. The code to reproduce these experiments is available at https://github.com/rgklab/StructuredNNs.
5.1 Density Estimation on Binary Data
For binary density estimation, we compare StrNN against the fully autoregressive MADE baseline.
Synthetic Tabular Data We generate binary tabular data through structural equations from known Bayesian networks. The results are shown in Figure 3 (left). We find that StrNN performs better than MADE, especially in the low data regime, as demonstrated on the left hand side of each chart.
MNIST Image Data To study the effect of structure in image modeling, we use the binarized MNIST dataset considered in Germain et al. [2015], Salakhutdinov and Murray [2008]. Germain et al. [2015] treated each 28-by-28-pixel image as a 784-dimensional data vector with full autoregressive dependence. Since we do not know the ground truth structure, we use StrNN to model a local autoregressive dependence on a square of a pixels determined by the hyper-parameter nbr_size. By changing the hyperparameter we can increase the context window used to model each pixel.
Figure 3: Negative log likelihood of a test set (lower is better) for density estimation experiments on binary synthetic data (left 2 images) and label-dependent binary MNIST data (right 2 images). Error ranges are 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_size=28. We first find the best nbr_size for each label via grid search. For labels 0 and 2, the optimal nbr_size is 10, and per-label density estimation results can be found in Figure 3 (right). StrNN outperforms MADE for both labels, with the advantage more significant when sample size is very small. Samples of handwritten digits generated from both StrNN and MADE can be found in Appendix B.
5.2 Density Estimation on Gaussian Data
We run experiments to compare the performances of StrNN and MADE on synthetic Gaussian data generated from known structure equation models where each  $x_{i}$  is Gaussian. We plot the results in Figure 4. StrNN achieves lower test loss than MADE on average, although the error bars are not necessarily disjoint. When the sample size is low, however, StrNN significantly outperforms MADE, similar to the binary case. In conclusion, across all binary and Gaussian experiments, encoding structure makes StrNN significantly more accurate at density estimation than the fully autoregressive MADE baseline.
Figure 4: Results from density estimation experiments on Gaussian synthetic data generated from two different 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 on average.
5.3 Density Estimation with Structured Normalizing Flows
We evaluate StrAF on density estimation against several baselines. We draw 1000 samples from a 15 dimensional tri-modal and non-linear synthetic dataset for experimental evaluation. Data generation is further described in Appendix C.4.
Experimental Setup We select the fully autoregressive flow (ARF) and the Graphical Normalizing Flow (GNF) [Wehenkel and Louppe, 2021] as the most relevant discrete flow baselines for comparison. We use the official GNF code repository during evaluation, but note that it contains design decisions that harm sample quality (see Appendix E.3). We also evaluate against FFJORD [Grathwohl et al., 2018] and Weilbach et al. [2020] as baselines for StrCNF. While other structured flows exist and have been examined in Section 4, they do not represent variables using an autoregressive structure and hence are less comparable. Where applicable, models were provided the true adjacency matrix in their conditioners. All discrete flow models use a UMNN [Wehenkel and Louppe, 2019] transformer and we grid-search other hyperparameters as described in Appendix E.2. We evaluate density estimation performance using the negative log-likelihood (NLL). After fixing the hyperparameters, we perform 8 randomly initialized training runs, then report the mean and  $95\%$  CI of the test NLL for these runs.
StrNN improves flow-based models for density estimation
We report results in Table 1 and observe several trends. For both discrete and continuous flows, the ability to incorporate structure yields performance benefits compared to the ARF and FFJORD baselines. In particular, the StrNN offers advantages in comparison to baseline approaches that can encode structure. For example, while the method proposed by Weilbach et al. [2020] allows a DAG structure to
Figure 5: Model generated samples are shown in blue dots for randomly selected dimensions. The ground truth density is visualized by the orange contours. See Appendix E.3 for an explanation on why GNF performs poorly.
be injected into the baseline CNF, its inability to easily use a multilayered neural network to represent dynamics hinders performance as compared to the StrCNF. We observe that StrAF and GNF perform comparably, and the StrCNF outperforms all other models. We visualize the quality of samples generated by these flow models in Figure 5, and find that both StrAF and StrCNF yields samples that closely match the ground truth distribution.
5.4 Causal Inference with Structured Autoregressive Flows
We conduct synthetic experiments where we generate data according to a linear additive SEM. Details on data generation can be found in Appendix D.2. In our experiments involving 5- and 10-variable
SEMs, we compare StrAF against CAREFL, which only utilizes MADE as the conditioner and maintains autoregressive ordering of the variables. This comparison highlights the additional benefits of enforcing the generative model to be faithful to the causal graph, which is a feature unique to StrAF. Furthermore, unlike the previous work conducted by Khemakhem et al. [2021] that evaluates causal queries on individual variables alone, we propose a comprehensive evaluation metric called total mean squared error (MSE) for these two causal tasks as outlined below. We report the mean errors along with the standard deviations from multiple training runs with different datasets. The formulations of the metrics can be found in Appendix D.3. In addition, Appendix D.1 provides detailed algorithms and related discussions on generating interventional samples and computing counterfactuals with flows.
Table 1: Evaluation of flow-based models. Mean and  $95\%$  CI of test NLL over 8 runs reported.
Test NLL (↓)
ARF	-3.09 ± 0.43
GNF	-3.63 ± 0.35
StrAF	-3.55 ± 0.20
FFJORD	-1.85 ± 0.64
Weilbach	-2.59 ± 0.58
StrCNF	-4.01 ± 0.12
(a) Interventions
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 and queries.
(b) Counterfactuals
Interventions. We perform interventions on each variable  $x_{j}$  within the SEM and draw samples from the intervened causal system. Interventions are performed by setting the intervened value  $\alpha$  using one of eight integers perturbed from the mean of the intervened variable. We compute the expectations  $\mathbb{E}[x_i|do(x_j = \alpha)]$  for variable  $x_{i}$ , excluding the intervened variable  $x_{j}$  itself and its
preceding variables as they remain unaffected. To estimate these expectations, we use 1000 samples and calculate the squared error for the predictions. The resulting error is then averaged over the set of intervened values, intervened variables, and the corresponding variables with computed expectations under interventions. We present the plot of this aggregated metric, referred to as the total intervention mean squared error (total I-MSE), for the two SEMs in Figure 6a. Notably, StrAF consistently outperforms CAREFL across all training dataset sizes, highlighting the effectiveness of the additional graph structure in improving StrAF's inference of the interventional distribution.
Counterfactuals. In a similar setting as the intervention experiments, we evaluate StrAF and CAREFL on their ability to compute accurate counterfactuals conditioned on observed data. Counterfactual inference tackles what-if scenarios: determining the value of variable  $x_{i}$  if variable  $x_{j}$  had taken a different value  $\alpha$ . Unlike interventions, counterfactual queries involve deriving the latent variables  $\mathbf{z}$  given the observed data  $\mathbf{x} = \mathbf{x}_{\mathrm{obs}}$ , rather than sampling new  $\mathbf{z}$ . We generate 1000 observations using the synthetic SEM and derive counterfactual values  $\mathbf{x}$  by posing queries with varying  $\alpha$  values for each variable  $x_{j}$ . Similar to interventions, we compute the squared error and average over the 1000 observed samples and all possible combinations of counterfactual queries for each observed sample, and we refer to this metric as the total counterfactual mean squared error (total C-MSE). Figure 6b illustrates that, for both SEMs, StrAF outperforms CAREFL in making more accurate counterfactual predictions. Moreover, StrAF demonstrates consistent performance even in scenarios with limited available samples.
We introduce the Structured Neural Network (StrNN), which enables us to encode functional invariances in arbitrary neural networks via weight masking during learning. For density estimation tasks where the true dependencies are expressed via Bayesian networks (or adjacency matrices), we show that StrNN outperforms a fully autoregressive MADE model on synthetic and MNIST data. We integrate StrNN in autoregressive and continuous flow models to improve both density estimation and sample quality. Finally, we show that our structured autoregressive flow-based causal model outperforms existing baselines on causal inference. We address some limitations and directions for future work below.
Access to true adjacency structure. We assume access to the true conditional independence structure. While this information is available in many contexts, such as from domain experts, there is also a wide literature on learning the conditional independencies directly from data [Drton and Maathuis, 2017]. One prominent example is the NO-TEARS algorithm [Zheng et al., 2018]. Wehenkel and Louppe [2021] shows that integrating NO-TEARS with an autoregressive flow can improve density estimation when a ground truth adjacency is unknown. Any adjacency matrix learned from data can also be integrated in StrAF. In our work, we have found two existing causal structure discovery libraries that are relatively comprehensive and easy to use: Kalainathan et al. [2020] and Zheng et al. [2023]. Future work will use StrNN to directly learn structure from data, providing a full pipeline from structure discovery to density estimation and sample generation.
StrNN optimization objectives. The mask factorization algorithm used by StrNN can maximize different objectives while ensuring the matrices satisfy a sparsity constraint. In Section 3.1, we proposed two such objectives and in Appendix A we demonstrated that they can impact model generalization. StrNN provides a framework with which it is possible to explore other objectives to impose desirable properties on neural network architectures. Moreover, investigating the effect of sparse structure on faster and easier training is a valuable direction [Frankle and Carbin, 2018]. For example, one approach would be to leverage connections between dropout and the lottery ticket hypothesis to randomly introduce sparsity into a neural network using StrNN weight masking.
We thank David Duvenaud, Tom Ginsberg, Vahid Balazadeh-Meresht, Phil Fradkin, and Michael Cooper for insightful discussions and draft reviewing. We thank anonymous reviewers for feedback that has greatly improved the work. This research was supported by an NSERC Discovery Award RGPIN-2022-04546, and a Canada CIFAR AI Chair. AC is supported by a DeepMind Fellowship and RS is supported by the Ontario Graduate Scholarship and the Ontario Institute for Cancer Research. Resources used in preparing this manuscript were provided in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute.
Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems, 33:21865-21877, 2020.
Wenhao Ding, Haohong Lin, Bo Li, and Ding Zhao. Causalaf: Causal autoregressive flow for safety-critical driving scenario generation. In Conference on Robot Learning, pages 812-823. PMLR, 2023.
Rongguang Wang, Pratik Chaudhari, and Christos Davatzikos. Harmonization with flow-based causal inference. In Medical Image Computing and Computer Assisted Intervention-MICCAI 2021: 24th International Conference, Strasbourg, France, September 27-October 1, 2021, Proceedings, Part III 24, pages 181-190. Springer, 2021.
Sourabh Balgi, Jose M Pena, and Adel Daoud. Personalized public policy analysis in social sciences using causal-graphical normalizing flows. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 11810-11818, 2022a.
Sourabh Balgi, Jose M Pena, and Adel Daoud. Counterfactual analysis of the impact of the IMF program on child poverty in the global-south region using causal-graphical normalizing flows. arXiv preprint arXiv:2202.09391, 2022b.
Ruslan Salakhutdinov and Iain Murray. On the quantitative analysis of deep belief networks. In International Conference on Machine Learning, 2008.
Mathias Drton and Marloes H Maathuis. Structure learning in graphical modeling. Annual Review of Statistics and Its Application, 4:365-393, 2017.
Xun Zheng, Bryon Aragam, Pradeep K Ravikumar, and Eric P Xing. Dags with no tears: Continuous optimization for structure learning. Advances in neural information processing systems, 31, 2018.
Diviyan Kalainathan, Olivier Goudet, and Ritik Dutta. Causal discovery toolbox: Uncovering causal relationships in python. The Journal of Machine Learning Research, 21(1):1406-1410, 2020.
Yujia Zheng, Biwei Huang, Wei Chen, Joseph Ramsey, Mingming Gong, Ruichu Cai, Shohei Shimizu, Peter Spirtes, and Kun Zhang. Causal-learn: Causal discovery in python. arXiv preprint arXiv:2307.16405, 2023.