SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Asic Chen $^{1*}$  Ruian Shi $^{1*}$  Xiang Gao $^{1}$  Ricardo Baptista $^{2\dagger}$  Rahul G. Krishnan $^{1\dagger}$
<sup>1</sup>University of Toronto, Vector Institute <sup>2</sup>California Institute of Technology
{asicchen, ruiashi, xgao, rahulgk}@cs.toronto.edu
Injecting structure into neural networks enables learning functions that satisfy invariances with respect to subsets of inputs. For instance, when learning generative models using neural networks, it is advantageous to encode the conditional independence structure of observed variables, often in the form of Bayesian networks. We propose the Structured Neural Network (StrNN), which injects structure through masking pathways in a neural network. The masks are designed via a novel relationship we explore between neural network architectures and binary matrix factorization, to ensure that the desired independencies are respected. We devise and study practical algorithms for this otherwise NP-hard design problem based on novel objectives that control the model architecture. We demonstrate the utility of StrNN in three applications: (1) binary and Gaussian density estimation with StrNN, (2) real-valued density estimation with Structured Autoregressive Flows (StrAFs), autoregressive normalizing flows that leverage StrNN as a conditioner, and (3) interventional and counterfactual analysis with StrAFs. Our work opens up new avenues for learning neural networks that enable data-efficient generative modeling and the use of normalizing flows for causal effect estimation.
The incorporation of structure into machine learning models has been shown to provide benefits for model generalization, learning efficiency, and interpretability. The improvements are particularly salient when learning from small amounts of data. This idea has found use in reinforcement learning [Ok et al., 2018], computational healthcare [Hussain et al., 2021, Cui et al., 2020], survival analysis [Gharari et al., 2023], time series analysis [Curi et al., 2020], and causal inference [Balazadeh et al., 2022].
This work focuses on the problem of density estimation from high-dimensional data which has been approached through a variety of lenses. For example, normalizing flows [Tabak and Turner, 2013, Rezende and Mohamed, 2015] model data by transforming a base distribution through a series of invertible transformations. Masked autoencoders (MADE) [Germain et al., 2015] model the joint distribution via an autoregressive factorization of
$$
p (X) = p _ {\theta_ {1}} (x _ {1}) p _ {\theta_ {2}} (x _ {2} | x _ {1}) p _ {\theta_ {3}} (x _ {3} | x _ {2}) p _ {\theta_ {4}} (x _ {4} | x _ {2}, x _ {3})
$$
Figure 1: StrNN injects structure by masking the weights of a neural network. Top: StrNN connections (green) compared to a fully connected network (gray). Bottom: Binary factorization of an adjacency matrix yields weight masks. Masked weights shown in gray.
*Equal Contribution † Equal Senior Authorship
the random variables. The factorization is enforced by integrating structure in the neural network of the autoencoder. The MADE architecture zeros out weights in a neural network to ensure each output dimension has an autoregressive dependence on the input dimensions. When data is scarce, this may lead to over-fitting and harm generalization. When knowledge of a Bayesian network [Pearl, 2011] and the associated conditional independencies exist, it is desirable to inject this knowledge directly into the network to improve density estimation. In this work, we extend the concept of weight masking, as in MADEs, to go beyond autoregressive dependencies of the output on the input.
This work proposes the Structured Neural Network (StrNN), a network architecture that enforces functional independence relationships between inputs and outputs via weight masking. In other words, the output can remain unaffected by changes to (subsets of) the input. We focus on instantiating this idea to model conditional independence between inputs when neural networks are deployed to model the density of random variables, as illustrated in Figure 1. Any set of conditional independence statements (e.g. in a Bayesian network) may be represented via a binary adjacency matrix. StrNN performs binary matrix factorization to generate a set of weight masks that follow the adjacency matrix. There are two key challenges we overcome. First, the general problem of binary matrix factorization in this context is under-specified, as there exist many valid masks whose matrix product realizes a given adjacency matrix. To this end, we propose the idea of neural network path maximization as a strategy to guide the generation of optimal masks. Secondly, binary matrix factorization is NP-hard in general. We study practical solutions that generate valid mask matrices efficiently.
StrNN can then be applied to NN-based density estimation in various contexts. Where conditional independence properties are known a priori, we show that StrNN can be used to estimate parameters of data distributions while keeping specified variables conditionally independent. We further integrate StrNNs into various discrete and continuous flow architectures, including the autoregressive normalizing flow [Papamakarios et al., 2017, Huang et al., 2018] to form the Structured Autoregressive Flow (StrAF). The StrAF model uses the StrNN as a normalizing flow conditioner network, thus enforcing a given adjacency structure within each flow layer. The StrAF preserves variable orderings between chained layers, allowing the adjacency structure to be respected throughout the entire flow.
Finally, we study a natural application of StrAF in causal effect estimation, a domain that often requires flexible methods for density estimation that enforce conditional independence relationships within causal graphs [Pearl, 2009]. We show how StrAFs can be used to perform interventional and counterfactual queries better than existing flow-based causal models that do not incorporate graphical structures. Across the board, we highlight how incorporating conditional independence structure improves generalization error when learning from a small number of samples.
1. We introduce StrNN, a weight-masked neural network that can efficiently learn functions with specific variable dependence structures. In particular, it can inject prior domain knowledge in the form of Bayesian networks to probability distributions. We formalize the weight masking as an optimization problem, where we can pick the objective based on desired neural architectures. We propose an efficient binary matrix factorization algorithm to mask arbitrary neural networks.
2. We integrate StrNN into autoregressive and continuous normalizing flows for best-in-class performance in density estimation and sample generation tasks.
3. We apply StrAF for causal effect estimation and showcase its ability to outperform existing causal flow models in accurately addressing interventional and counterfactual queries.
Masked Autoencoders for Density Estimation (MADE): Masked neural networks were introduced for density estimation on binary-valued data [Germain et al., 2015]. Given  $\mathbf{x} = (x_{1},\dots,x_{d})$ , MADE factorizes  $p(\mathbf{x})$  as the product of the outputs of a neural network. Writing the  $j$ -th output as the conditional probability  $\hat{x}_j \coloneqq p(x_j = 1|\mathbf{x}_{< j})$ , the joint distribution can be rewritten exactly as the binary cross-entropy loss. As long as the neural network outputs are autoregressive in relation to its inputs, we can minimize the cross-entropy loss for density estimation. To enforce the autoregressive property for a neural network  $y = f(x)$  with a single hidden layer and  $d$  inputs and outputs, MADE element-wise multiplies the weight matrices  $\mathrm{W}$  and  $\mathrm{V}$  with binary masks  $M^W$  and  $M^V$ :
$$
h (x) = g \left(\left(W \odot M ^ {W}\right) x + b\right), \quad y = f \left(\left(V \odot M ^ {V}\right) h (x) + c\right). \tag {1}
$$
The autoregressive property is satisfied as long as the product of the masks,  $M^V M^W \in \mathbb{R}^{d \times d}$ , is lower triangular. The MADE masking algorithm (2) can be extended to neural networks with an arbitrary number of hidden layers and hidden sizes. For Gaussian data, the MADE model can be extended as  $\mathbb{R}^d \to \mathbb{R}^{2d}$ ,  $(x_1, \ldots, x_d) \to (\hat{\mu}_1, \ldots, \hat{\mu}_d, \log(\hat{\sigma}_1), \ldots, \log(\hat{\sigma}_d))$ . The last mask must be duplicated to ensure  $\mu_j$  and  $\sigma_j$  only depend on  $\mathbf{x}_{<j}$ . MADE can also be used as the conditioner in an autoregressive flow to model general data, as seen in Papamakarios et al. [2017].
Normalizing Flows: Normalizing flows [Rezende and Mohamed, 2015] model complex data distributions and have been applied in many scenarios [Papamakarios et al., 2021]. Given data  $\mathbf{x} \in \mathbb{R}^d$ , a normalizing flow  $\mathbf{T} \colon \mathbb{R}^d \to \mathbb{R}^d$  takes  $\mathbf{x}$  to latent variables  $\mathbf{z} \in \mathbb{R}^d$  that are distributed according to a simple base distribution  $p_{\mathbf{z}}$ , such as the standard normal. The transformation  $\mathbf{T}$  must be a diffeomorphism (i.e., differentiable and invertible) so that we can compute the density of  $\mathbf{x}$  via the change-of-variables formula:  $p_{\mathbf{x}}(\mathbf{x}) = p_{\mathbf{z}}(\mathbf{T}(\mathbf{x}))|\det J_{\mathbf{T}}(\mathbf{x})|$ . We can compose multiple diffeomorphic transformations  $\mathbf{T}_k$  to form the flow  $\mathbf{T} = \mathbf{T}_1 \circ \dots \circ \mathbf{T}_K$  since diffeomorphisms are closed under composition. The flows are trained by maximizing the log-likelihood of the observed data under the density  $p_{\mathbf{x}}(\mathbf{x})$ . The log-likelihood can be evaluated efficiently when it is tractable to compute the Jacobian determinant of  $\mathbf{T}$ ; for example when  $\mathbf{T}_k$  is a lower triangular function [Marzouk et al., 2016]. Given the map, we can easily generate i.i.d. samples from the learned distribution by sampling from the base distribution  $\mathbf{z}^i \sim p_{\mathbf{z}}$  and evaluating the flow  $\mathbf{T}^{-1}(\mathbf{z}^i)$ .
Density Estimation with Autoregressive Flows: When the Jacobian matrix of each flow layer is lower triangular, its determinant is simply the product of its diagonal entries [Huang et al., 2018]. This gives rise to the autoregressive flow formulation: given an ordering  $\pi$  of the  $d$  variables in the data vector  $\mathbf{x}$ , the  $j$ th component of the flow  $\mathbf{T}$  has the form:  $x_{j} = \tau_{j}(z_{j};c_{j}(\mathbf{x}_{<\pi(j)}))$  where each  $\tau_{j}$  is an invertible transformer and each  $c_{j}$  is a conditioner that only depends on the variables that come before  $x_{j}$  in the ordering  $\pi$ . As a result, the map components define an autoregressive model that factors the density over a random variable  $x$  as:  $p(\mathbf{x}) = \prod_{j=1}^{d} p(x_{j}|\mathbf{x}_{<j})$  where  $\mathbf{x}_{<j} = (x_{1},\ldots,x_{j-1})$ . Autoregressive. Under mild conditions, any arbitrary distribution  $p_{\mathbf{x}}$  can be transformed into a base distribution with a lower triangular Jacobian [Rezende and Mohamed, 2015]. That is, autoregressive flows are arbitrarily expressive given the target distribution. One common choice of invertible functions for the transformer are monotonic neural networks [Wehenkel and Louppe, 2019].
Density Estimation with Continuous Normalizing Flows: Continuous normalizing flows (CNFs) [Chen et al., 2018, Grathwohl et al., 2018] represent the transformation  $\mathbf{T}$  as the flow map solving the differential equation  $\frac{\partial\mathbf{z}(t)}{\partial t} = f(\mathbf{z}(t),t;\theta)$ . Given the initial condition  $\mathbf{x} = \mathbf{z}(t_1)\sim p_{\mathbf{x}}$ , we can integrate  $f$  backwards in time from  $t_1$  to  $t_0$  to obtain  $\mathbf{z} = \mathbf{z}(t_0)\sim p_{\mathbf{z}}$ , or vice versa. In order to learn the CNF, Chen et al. [2018] computes the change in log-density under the transformation using the instantaneous change of variables formula, which is defined by the differential equation  $\frac{\partial\log p(\mathbf{z}(t))}{\partial t} = -\mathrm{Tr}(\frac{\partial f}{\partial\mathbf{z}(t)})$ . This expression is used to compute the log-likelihood of a target sample as  $\log p(\mathbf{z}(t_1)) = \log p(\mathbf{z}(t_0)) - \int_{t_0}^{t_1}\mathrm{Tr}(\frac{\partial f}{\partial\mathbf{z}(t)})dt$ . We refer the reader to Chen et al. [2018] for the process of back-propagating through the objective using the adjoint sensitivity method as well as a discussion on the existence and uniqueness of a solution for the ODE flow map.
Causal Inference with Autoregressive Flows: Modelling causal relationships is crucial for enabling effective decision-making in various fields [Pearl, 2009]. A structural equation model (SEM) parameterizes the process that generates observed data, allowing us to reason about interventions and counterfactuals. Given random variables  $\mathbf{x} = (x_{1},\dots,x_{d})\in \mathbb{R}^{d}$  with joint distribution  $\mathbb{P}_{\mathbf{x}}$ , the associated SEM consists of  $d$  structural equations of the form  $x_{j} = f_{j}(\mathbf{pa}_{j},u_{j})$ , where  $u$  represents mutually independent latent variables and  $\mathbf{pa}_j$  denotes the direct causal parents of variable  $x_{j}$ . Each SEM also corresponds to a directed acyclic graph (DAG), with a causal ordering  $\pi$  defined by the variables' dependencies.
Khemakhem et al. [2021] showcased the intrinsic connection between SEMs and autoregressive flows. The authors demonstrated that affine autoregressive flows with a fixed ordering of variables can be used to parameterize SEMs under the framework of causal autoregressive flows (CAREFL). When the causal ordering of variables is given, CAREFL outperforms other baselines in interventional tasks and generates accurate counterfactual samples. However, one significant limitation of CAREFL is that there is no guarantee that the autoregressive structure corresponds to the true dependencies in the causal graph beyond pairwise examples. In Section 3.3, we leverage StrAF to incorporate additional independence structure, enhancing its faithfulness to the causal DAG.
3.1 Structured Neural Networks
As neural networks are universal function approximators [Hornik et al., 1989], injecting structure into multilayer perceptrons (MLPs) naturally allows us to approximate complex functions with certain invariances. For a neural network that approximates an arbitrary function  $f: \mathbb{R}^m \to \mathbb{R}^n$ , one type of invariance we could consider is when a specific output  $\hat{x}_j$  is independent from a given input  $x_i$ , i.e.:  $\hat{x}_j \perp x_i$  means  $\frac{\partial \hat{x}_j}{\partial x_i} = 0$ . As a motivating example, we focus on the probabilistic density estimation problem framed as learning maps from one probability distribution to another. In this setting, it is imperative that we are able to learn structured functions between two high-dimensional spaces. Therefore in this paper, we seek to efficiently encode such invariances via weight masking.
For data  $\mathbf{x} = (x_{1},\dots,x_{d})$ , we use the lower-triangular adjacency matrix  $A\in \{0,1\}^{d\times d}$  to represent the underlying variable dependence structure. In other words,  $A_{ij} = 0$  for  $j < i$  if and only if  $x_{i}\perp x_{j}|x_{\{1,\ldots ,i\} \setminus j}$  and  $A_{ij} = 1$  otherwise. This matrix encodes the same information as a Bayesian network DAG of the variables. In the fully autoregressive case, matrix A is a dense lower triangular matrix with all ones under the diagonal, which is the only case addressed in Germain et al. [2015]. Their proposed MADE algorithm (Appendix A.1) only encodes the structure of dense adjacency matrices, and cannot incorporate additional conditional independencies. Further, the non-deterministic version of the algorithm can introduce unwanted independence statements by chance, as discussed in Proposition 1.
We improve upon the idea of masked autoregressive neural networks to directly encode the independence structure represented by an adjacency matrix  $A$  that is lower triangular but also has added sparsity. We observe that a masked neural network satisfies the structural constraints prescribed in  $A$  if the product of the masks for each hidden layer has the same locations of zero and non-zero entries as  $A$ . Therefore, given the conditional independence structure of the underlying data generating process, we can encode structure into an autoregressive neural network by factoring the adjacency matrix into binary mask matrices for each hidden layer.
More specifically, given an adjacency matrix  $A \in \{0,1\}^{d \times d}$  and a neural network with  $L$  hidden layers, each with  $h_1, h_2, \ldots, h_L$  hidden units ( $\geq d$ ), we seek mask matrices  $M^1 \in \{0,1\}^{h_1 \times d}$ ,  $M^2 \in \{0,1\}^{h_2 \times h_1}$ , ...,  $M^L \in \{0,1\}^{d \times h_L}$  such that  $A' \sim A$ , where  $A' := M^L \cdots M^2 \cdots M^1$ . We use  $A' \sim A$  to denote that matrices  $A'$  and  $A$  share the same sparsity pattern, i.e.: exact same locations of zeros and non-zeros. Note that here  $A$  is a binary matrix and  $A'$  is an integer matrix. We then mask the neural network's hidden layers using  $M^1, M^2, \ldots, M^L$  as per (1) to obtain a Structured Neural Network (StrNN), which respects the prescribed independence constraints. The value of each entry  $A_{ij}'$  thus corresponds to the number connections flowing from input  $x_j$  to output  $\hat{x}_i$  in the StrNN.
Finding the optimal solution to this problem is NP-hard since binary matrix factorization can be reduced to the biclique covering problem [Miettinen and Neumann, 2020, Ravanbakhsh et al., 2016, Orlin, 1977]. Furthermore, most existing works focus on deconstructing a given matrix  $A$  into low-rank factors while minimizing (but not eliminating) reconstruction error [Dan et al., 2015, Fomin et al., 2020]. In our application, any non-zero reconstruction error breaks the independence structure we want to enforce in our masked neural network. This puts existing algorithms for low-rank binary matrix factorization outside the scope of our paper. We instead consider the problem of finding factors that reproduce the adjacency matrix exactly, which is always possible when hidden layer dimensions are greater than the input and output dimension.
Optimization Objectives. Identifiability remains an issue even when we eliminate reconstruction error. Given an adjacency matrix  $A$ , there can be multiple solutions for factoring  $A$  into per-layer masks that satisfy the constraints, especially if the dimensions of the hidden layers are much larger than  $d$ . Since the masks dictate which connections remain in the neural network, the chosen mask factorization algorithm directly impacts the neural network architecture. Hence, it is natural to explicitly specify a relevant objective to the neural network's approximation error during the matrix factorization step, such as the test log-likelihood.
Given that the approximation error is inaccessible when selecting the architecture, we were inspired by the Lottery Ticket Hypothesis [Frankle and Carbin, 2018] and other pruning strategies [Srivastava et al., 2014, Gal and Ghahramani, 2016] that identify a subset of valuable model connections. Our hypothesis is that given the same data and prior knowledge on independence structure, the masked
neural network with more connections is more expressive, and will thus be able to learn the data better and/or more quickly. To find such models, we consider two objectives: Equation (2) that maximizes the number of connections in the neural network while respecting the conditional independence statements dictated by the adjacency matrix, and Equation (3) that maximizes connections while penalizing any pair of variables from having too many connections at the cost of the others. That is,
$$
\max  _ {A ^ {\prime} \sim A} \sum_ {i = 1} ^ {d} \sum_ {j = 1} ^ {d} A _ {i j} ^ {\prime}, \tag {2}
$$
$$
\max  _ {A ^ {\prime} \sim A} \sum_ {i = 1} ^ {d} \sum_ {j = 1} ^ {i} A _ {i j} ^ {\prime} - \operatorname {v a r} \left(A ^ {\prime}\right), \tag {3}
$$
where  $\operatorname{var}(A')$  is the variance across all entries in  $A'$ . While we focus on these two objectives, future work will find optimal architectures by identifying other objectives to improve approximation error.
Factorization Algorithms. The maximization of the aforementioned mask factorization is an intractable optimization problem. We therefore develop approximate algorithms to solve them. Our strategy for optimization involves recursively factorizing the mask matrix layer by layer. Given  $A \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 that the matrices share the same sparsity. For the next layer with  $h_2$  hidden units, we use  $A_1$  in place of  $A$  to find  $A_2 \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$ . We repeat until we have found all the masks.
For each objective, we can obtain per-layer exact solutions using integer programming. While the Gurobi optimizer [Gurobi Optimization, LLC, 2023] can be used for small  $d$ , this approach was found to be too computationally expensive for  $d$  greater than 20, which is a severe limitation for real-world datasets and models. We hereby propose a greedy algorithm (shown in Algorithm 1) that approximates the solution to the maximum connections objective in Equation (2). For each layer, the algorithm first replicates the structure given in the adjacency matrix  $A$  by copying its rows into the first mask. It then maximizes the number of neural network connections by filling in the second mask with as many ones as possible while respecting the sparsity in  $A$ . See Appendix A.2 for a visual explanation of the algorithm. For a network with  $d$ -dimensional inputs and outputs and one hidden layer with  $h$  units, this algorithm runs in  $\mathcal{O}(dh)$  time. Scaling up to  $L$  layers, where each hidden layer commonly contains  $\mathcal{O}(d)$  units, the overall runtime is  $\mathcal{O}(d^2 L)$ , which is much more efficient than the integer programming solutions. From our experiments, the greedy algorithm executes nearly instantaneously for dimensions in the thousands.
In Appendix A, we include detailed results from investigating the link between the neural network's generalization performance and the choice of mask factorization algorithm. We observe that while the exact solution to objective (2) achieves a higher objective value than the greedy approach, it has no clear advantage in density estimation performance. Moreover, we found that models trained with the two objectives, (2) and (3), provide similar performance. However, some datasets might be more sensitive to the exact objective. For example, problems with anisotropic 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-
Algorithm 1: Greedy factorization
$$
M ^ {W} \in \{0, 1 \} ^ {h \times d _ {2}}, \text {s a t i s f y i n g}
$$
$$
M ^ {V} M ^ {W} \sim A
$$
Data:  $A\in \{0,1\}^{d_1\times d_2}$  , hidden size h
Result:  $M^{\vee}\in \{0,1\}^{d_1\times h},$
1  $nz\gets$  non-zero rows in  $A$
2 Fill  $M^W$  with  $nz$ ; repeat until full
3 Fill  $M^V$  with ones
4 for  $i$ -th row in  $M^V$  do:
5  $C\gets$  indices of 0's in  $i$  -th row of  $A$
$T\gets$  cols.of  $M^W$  whose index  $\in C$
7  $R\gets$  indices of non-zero rows of  $T$
for  $r$  in  $R$ : set  $M_{i,r}^{V}$  to zero
9 return  $(M^V,M^{W})$
ing different factorization objectives is an important direction for future work.
3.2 Structured Neural Networks for Normalizing Flows