Structured Neural Networks for Density Estimation and Causal Inference

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.

Paper

References (52)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer LoZH6/10 · confidence 3/52023-06-26

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. **Edit** Most of my concerns were addressed during the rebuttal/discussion period. I am therefore upgrading my score from 4 to 6.

Strengths

* The manuscript is well written, with standard notations. * The contributions are clearly defined and the assumptions (known adjacency) are explicit. * The proposed approach to impose conditional independences is sound. * 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: > 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). In 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.

Questions

* Algorithm 1: It is not clear to me how the factorization algorithm is applied when the StrNN has more than one hidden layer. * Section 5.1: Are the same number of layers/neurons used for StrNN and MADE in this experiment? * Line 317: "As GNF permutes variables between flow steps". I was not able to find a mention of this in [25]. * 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"? * 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. * 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

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. * 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.

Reviewer icwN6/10 · confidence 4/52023-06-29

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.

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. 2. To tackle the infeasibility of exact mask creation on larger graphs a greedy algorithm is proposed and its practical application is demonstrated. 3. 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.

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. 2. 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. Overall 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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

No concerns here

Reviewer Wq9i6/10 · confidence 4/52023-07-06

Summary

In this paper, the authors present a neural network architecture that can fulfill the bayesian DAG conditional independencies needed for normalized density estimation. In 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. The authors introduce an heuristic to exactly factorize the adjacency matrix for the different layers using two objective functions. With 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.

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. The 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. The 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. A broader comparison on other datasets would make the paper more robust. Also 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.

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. 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.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. At 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. I'm wondering if you explored similar ideas during your research?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

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. There are no potential negative societal impact to this work.

Reviewer Tj4q5/10 · confidence 4/52023-07-27

Summary

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

Strengths

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

Weaknesses

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

Questions

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

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Reviewer icwN2023-08-11

Thank you for the Response

I would like to thank the authors for their response. I am happy with their response and am still positive about the paper and it's contributions.

Reviewer LoZH2023-08-11

Thank you for your answers and the experiments you have conducted in such a short time. You have addressed most of my concerns. I appreciate your honesty regarding the similarities between your factorization algorithm(s) and the one present in Zuko. I also appreciate you taking the time to understand and highlight the differences, which, as you demonstrate, can impact the performances beyond enforcing the independences. I would like to mention that, to my understanding, one factorization is not necessarily better than the other. They are different and will lead to different results for different adjacency structures. Nevertheless, you rightfully observe that, even though the factorization algorithm is present in Zuko, users are not easily able to initialize their flow to respect a particular adjacency matrix. After some research, I did not find a publication linked to the Zuko library, or to the factorization algorithm therein. I believe a reference to the repo as well as a short discussion of the similarities/differences would be enough. Finally, I agree with you that there is a gap between the suggestion of an idea and its implementation, and your work goes beyond this idea by demonstrating the impact of the factorization in terms of generalization. Concerning my questions, I thank you for your clear answers. I would suggest you to add some of these within the manuscript/discussion. In view of the rebuttal and the proposed changes, I am willing to reconsider my score. I believe your work, although light from a theory perspective, is a valuable contribution to the density estimation community.

Authorsrebuttal2023-08-11

Dear Reviewer LoZH: Thank you for the prompt response to our rebuttal. We are very glad that we have managed to address most of your concerns. We agree it is unlikely there is one factorization scheme which is optimal for all possible adjacency matrices, which is a limitation we will discuss in our revision. However, we do believe some factorization objectives/algorithms are better than others on average due to how they dictate the resulting masked NN architecture. We have identified this as an important area to formalize and investigate in future work. In our rebuttal, we used one example where our greedy algorithm outperforms Zuko to highlight one key difference. At the same time, we acknowledge that there might be specific adjacency structures for which our greedy algorithm would do less well. Nonetheless, StrNN also gives us the ability to choose our factorization objective based on prior information, which we believe will help shed light on this in the future. We are currently editing our manuscript for clarity based on all reviewer suggestions and questions. Thank you again for helping us improve our work! If you do find our responses satisfactory, we would greatly appreciate it if you could re-evaluate our paper for a higher score. Thank you for your time. Best, Authors

Reviewer Wq9i2023-08-14

I thank the authors for their rebuttal. I still consider the paper in a positive light and after following the discussion, I will keep my score as is.

Program Chairsdecision2023-09-21

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC