Flow-based parameterization for DAG and feature discovery in scientific multimodal data

Abstract

Representation learning algorithms are often used to extract essential features from high-dimensional datasets. These algorithms commonly assume that such features are independent. However, multimodal datasets containing complementary information often have causally related features. Consequently, there is a need to discover features purporting conditional independencies. Bayesian networks (BNs) are probabilistic graphical models that use directed acyclic graphs (DAGs) to encode the conditional independencies of a joint distribution. To discover features and their conditional independence structure, we develop pimaDAG, a variational autoencoder framework that learns features from multimodal datasets, possibly with known physics constraints, and a BN describing the feature distribution. Our algorithm introduces a new DAG parameterization, which we use to learn a BN simultaneously with a latent space of a variational autoencoder in an end-to-end differentiable framework via a single, tractable evidence lower bound loss function. We place a Gaussian mixture prior on the latent space and identify each of the Gaussians with an outcome of the DAG nodes; this identification enables feature discovery with conditional independence relationships obeying the Markov factorization property. Tested against a synthetic and a scientific dataset, our results demonstrate the capability of learning a BN on simultaneously discovered key features in a fully unsupervised setting.

1 Introduction

To achieve autonomous scientific discovery, scientists are rapidly collecting large scientific datasets with a growing number of complex modalities. Although traditional multimodal datasets may consist of analogous text, image, and video modalities, these scientific datasets may contain disparate modalities with varying fidelity and information, such as 0D process parameters, 3D scanning electron microscopy imagery, and 1D X-ray fluorescence spectroscopy. Such large, multimodal scientific datasets extend beyond the limits of human cognition and thereby necessitate machine learning (ML)-driven methods to identify hidden, underlying factors in the data (Boyce and Uchic, 2019; Sparkes et al., 2010). Machine learning methods for these tasks are increasingly being asked to perform multiple tasks at once: to discover hidden relationships that are latent to data, to fuse different data measurements and modalities for novel scientific insights, and to meaningfully link relationships in a causal manner. In this paper, we propose a novel method called pimaDAG, designed to encode multimodal data in a shared latent space, while simultaneously learning an underlying causal structure within the latent representation. We do so by enforcing a structured Gaussian mixture prior on the latent space and then discovering a Bayesian network that defines the mixing probabilities in the Gaussian mixture.

As our method seeks to perform multimodal latent representation learning and causal discovery simultaneously, we summarize common relevant approaches for each below, to situate our work in the context of existing methods.

1.1 Multimodal latent representation learning

The field of disentangled representation learning seeks to identify hidden features of data through an interpretable latent representation (Bengio et al., 2013). Variational autoencoder (VAE) frameworks are often used in representation learning to provide a meaningful, disentangled representation of data in a latent space (Higgins et al., 2017; Kingma and Welling, 2019). The original VAE latent space assumes that data follow a standard Gaussian prior once embedded into the latent space, tasking the encoders and decoders to transform the data’s natural distribution to match the latent normal prior; many other distributions for priors have been proposed since.

Multimodal VAE approaches (e.g., Khattar et al., 2019) additionally task the encoders and decoders to correlate features between modalities, imposing an implicit structure on the latent space in the sense that latent embedding must synthesize representative features for each modality. These additional modalities can improve classification and disentanglement (Walker et al., 2024). These modalities, which include scalar-valued data, time-series information, audio, and video, are ideally complementary, each present novel information about the same datum. This capability becomes increasingly important for scientific datasets, which additionally consist of various modalities and obey physics constraints.

In particular, for scientific tasks, physics-informed multimodal autoencoders (PIMAs) have demonstrated the ability to detect features in multimodal datasets while incorporating known physics to aid in disentanglement Walker et al., 2024. This approach uses a VAE framework to learn a joint representation of multimodal data with optional physical constraints on the decoders, fusing information from each modality through a product-of-experts model.

Occasionally, multimodal representation methods consider causal or correlative relationships, such as the method of Lyu et al. (2022), which aims to maximize latent cross-view correlations. In particular, PIMA and many other VAE methods do not consider any dependencies, including causal relationships, between its discovered features. For example, many VAE frameworks assume that features (and modalities) are independent. Real-world data, however, have natural correlative and causal relationships, which, in the context of this work, we wish to exploit so that we may interpret relevant underlying factors within the latent representations that are uncovered.

1.2 Causal discovery

As our aim is to discover latent features with plausible causal relationships, we provide a few snapshots into the field of causal learning. Our treatment is by no means a complete view of this field.

1.2.1 Causal representation learning

The general research area of causal representation learning focuses on identifying causal relationships among latent variables. One recurring goal of this relatively new, broad research area is identifying latent features from data and a Bayesian network relating those features (Schölkopf et al., 2021). Bayesian networks (BNs) model conditional dependencies within a joint distribution of random variables via directed acyclic graphs (DAGs) (Jensen, 2001). The set of nodes in the DAG represent the random variables of the joint distribution, and the DAG dictates conditional independencies of the variables by way of the Markov factorization property.

Efforts in causal representation learning are fairly broad, and many efforts exploit prior knowledge in order to identify a unique graph within a larger causal learning framework. Some common assumptions include linear structural models (Squires et al., 2023; Kocaoglu et al., 2018), oracles aiding the causal discovery process (Yang et al., 2021; Shen et al., 2022; Yao et al., 2024), time-dependent data (Lippe et al., 2022, 2023b,a; Löwe et al., 2022; Yao et al., 2022), or access to interventions (Yang et al., 2021; Squires et al., 2023; Buchholz et al., 2023; Varici et al., 2023).

A few of these methods consider multimodal data. For example, Yao et al. (2024) focused on identifiable multi-view learning when the graph of the latent variables is already known. As we decode to different modalities from a joint latent space, we operate under a similar multi-view assumption, but we instead recover our DAG by training on data, and our learned feature distribution is a Bayesian network. The methods presented in Morioka and Hyvarinen (2023) provide another multimodal algorithm, where the authors simultaneously perform causal discovery and representation learning. Unlike our method, they base their feature representation scheme on nonlinear independent component analysis (NICA; Hyvarinen and Morioka (2017)). Our scheme instead discovers a probabilistic shared representation of features, enabling us to robustly capture uncertain observations of multimodal features and provide for a larger range of representations by using arbitrary deep architectures for the encoders and decoders in our VAE.

Other works focus on temporal sequences to achieve identifiable causal representations. For example, Lippe et al. (2022, 2023b,a), Löwe et al. (2022), and Yao et al. (2022) build graphical causal representations, but they do so with time-series data using notions of Granger causality, which are distinct from Bayesian network discovery in the sense that time directionality complicates the notion of independence; these papers leverage temporal dynamics (or in the case of Löwe et al. (2022), amortize across multiple instances of the same dynamics) to build latent representations, from which these authors build notions of causality. In our case, we do not presume time-dependent data, and instead, we leverage common representations built from multimodal data to inform our causal discovery process. Consequently, our work is more general since we make no assumptions on the influence of time-dependence in the data.

1.2.2 Observational DAG discovery

In the interest of extracting unique DAGs, the important methods highlighted in Section 1.2.1 tend to make assumptions that are not feasible in real-world or scientific settings. For example, intervention data are often not practical or sometimes not possible to obtain, particularly in observational studies. Consider, for instance, a material science discovery setting where the process parameters used to create a sample are typically believed to affect the material’s microstructure, which, in turn, would affect the material’s properties and performance. In such a material science setting, data are typically collected on the process parameters, microstructure, and properties. However, there is no way to intervene in a way that the microstructure no longer affects the properties, meaning that perfect interventions are not possible. Furthermore, the expense of generating materials samples and data is not trivial, which might limit the ability to collect more than one sample for a given set of process parameters. Despite these challenges, however, there is still a need to discover a feature distribution in high-dimensional scientific data that have a conditional independence structure. As a result, observational causal discovery has drawn much attention in recent years (Glymour et al., 2019; Wang et al., 2024), especially in scientific settings. Bayesian networks with discovered features can help scientists formulate hypotheses to further direct their research. Bayesian networks on discovered features can help identify plausible causal relationships suitable for future investigation. Consequently, we pursue an unsupervised representation learning algorithm that learns a Bayesian network of discovered features from any general set of (high-dimensional, multimodal, and scientific) data. Specifically, our focus is thus to learn a DAG and corresponding joint distribution of features (i.e., a Bayesian network) in a smooth manner amenable to gradient descent and coupling with a scientific, multimodal VAE (e.g., PIMA).

DAG discovery for Bayesian networks is, in general, an NP-hard problem (Chickering et al., 2004), and as a result, algorithms for DAG discovery have been the subject of their own line of research. This literature encompasses a wide range of algorithms; score-based methods, conditional independence testing, and continuous optimization approaches are the most popular. Both conditional independence testing and score-based methods often rely upon combinatorial searches to test conditional independence or to otherwise enforce acyclicity of the learned DAG. Although recent algorithms, e.g., Ramsey et al. (2017), Shimizu et al. (2006), and Spirtes et al. (2000), have made strides in optimizing this search, these combinatorial methods still remain expensive. Recently, continuous, differentiable optimization methods have expedited the discovery of DAGs through equality conditions enforcing the acyclic constraint, thereby bypassing the otherwise laborious search in the space of all DAGs (Zheng et al., 2018; Wei et al., 2020).

A breakthrough for continuous optimization schemes for learning DAGs, entitled NO TEARS, was introduced in Zheng et al. (2018), which developed new conditions for enforcing acyclicity in directed graphs by reformulating the combinatorial graph problem to a nonconvex optimization problem with an equality constraint. Works such as Wei et al. (2020),Kalainathan et al. (2022),Ng et al. (2019), and Lee et al. (2020) further build off of this idea and introduce alternative continuous constraints for learning DAGs. Applications of continuous optimization of DAGs include Yu et al. (2021), Yu et al. (2019),Yang et al. (2021),Pamfil et al. (2020), and Gao et al. (2022). In contrast, our DAG parameterization is not constraint-based or penalty-based, but it rather is natural parameterization for DAGs inspired by the Hodge theory (Jiang et al., 2011; Lim, 2020), where we view edges as the flow of information between nodes. In our words, this means our parameterization defines a DAG exactly at every step of training. It is important to note that sub-steps in Zheng et al. (2018) (and related works) are not DAGs; only the final optimal solution of the continuous optimization problem yields a DAG. When performing both multimodal learning and causal discovery simultaneously, it is desirable that at every training step, we maintain a DAG structure, making such methods like the ones above less desirable. Furthermore, our novel parameterization includes a temperature parameter that regularizes the edge indicator function in order to avoid local minima while training while still maintaining a valid DAG throughout the continuous optimization training procedure. For more references to DAG learning, see Vowels et al. (2022).

1.3 Our method: pimaDAG

Building off the feature discovery of PIMA (Walker et al., 2024), we present pimaDAG, which simultaneously learns an efficient representation of multimodal data in a shared latent space of a VAE while also discovering causal structures between the features of the discovered latent representations. The unique contributions of pimaDAG are (1) a new DAG and BN parametrization and (2) the linking of the trainable BN to PIMA, resulting in simultaneous BN and feature discovery in multimodal data. To link PIMA with a trainable BN, we assume that a Gaussian mixture (GMM) prior structures the latent space and that the mixing probabilities of the clusters of that GMM are defined by a trainable Bayesian network. To train our BN simultaneously with the VAE, we pose a novel continuous optimization scheme for parameterizing DAGs to match data so as to construct an end-to-end training pipeline that performs DAG discovery as a sub-step of training a VAE. Such training must also be handled differently: the commonly used evidence-based lower bound (ELBO) loss for VAEs must be adapted to include the new DAG parameterization. The ELBO above is computationally tractable through strategic framework decisions: our framework (1) utilizes unimodal deep encodings with Gaussian outputs, (2) fuses the unimodal deep encodings via a product of experts (PoE), (3) models clusters in the latent space as a mixture of Gaussians, (4) computes the probability of each cluster as the joint probability of the nodes of a trainable DAG, and (5) utilizes a mixture of deep decoders with the optional capability of physics-informed decoders for modalities suitable for expert modeling. To force better clustering in the latent space, we adapt the expectation maximization (EM) algorithm for fitting Gaussian mixture models to data as a training sub-step while minimizing the ELBO in order to calibrate the Gaussian mixture prior over the course of training. The problem formulation and DAG parameterization are presented in Section 2.1, and the training scheme is presented in Section 2.2.

2 Methods

For pimaDAG, we make very few assumptions regarding the incoming datasets. Datasets can consist of one or many modalities, where the modalities can vary greatly in complexity, e.g., from scalars to images. Furthermore, the modalities may share common data or be fairly complementary, or otherwise, they may satisfy exploitable physical constraints. Regardless of the dataset, the goal is to find a joint representation of all the modalities, where key features have dependencies that are described by a DAG. To this end, we do assume that there is a set of discrete, discoverable features within the data, and we further assume that each feature is a discrete random variable in a DAG. The number of features (nodes) and the number of outcomes of each feature serve as hyperparameters in pimaDAG. Knowing the number of features and outcomes a priori would expedite training, but otherwise, one can perform a hyperparameter sweep on these values. We outline our framework in Section 2.1, where the variational autoencoder setup (see orange box (a) in Figure 1) is given in Section 2.1.1, and new DAG parameterization and nodal distribution parameterization are given in Section 2.1.2 (see purple box (b) in Figure 1). Section 2.2 contains training information, including the single-sample ELBO and practical considerations.

FIGURE 1

For convenience, Table 1 in summarizes the notation used throughout this section and the rest of the paper.

2.1 Algorithmic framework

Given data from distinct modalities, we seek a common embedding into a latent space , where the latent space representation admits distinct clusters based on encoded features of the data. To model these clusters, we assume that is a Gaussian mixture model (GMM). We identify our encoded GMM clusters with discrete features . The probability distribution of forms the mixing probabilities of the GMM. In other words, if each feature is a categorical random variable with outcomes and we let denote one outcome of , then is normally distributed and has the probability of belonging. Moreover, we assume that there are shared latent dependencies between features and that these dependencies are described by a directed acyclic graph (DAG): each feature is a node in a DAG . Thus, by assuming the Markov factorization property,where denotes the immediate parents of feature node . Supplementary Figure S1 shows how the DAG of features relates to the latent space embedding.

For such a Gaussian mixture latent representation described by and , we aim to construct our embedding and latent space representation through training a multimodal variational autoencoder. Following works such as Dilokthanakul et al. (2016), Jiang et al. (2017), and Walker et al. (2024), we train our VAE to find the prior distribution , the posterior distribution , the joint distribution of , and the DAG , which maximize the evidence lower bound (ELBO) loss:We assume independence of decoding mechanisms for each modality for our prior, and we assume mean-field separability for the posterior. These assumptions, respectively, giveWith these assumptions, the ELBO separates as sums of expectations of Gaussian distributions; see Equation 23. In the case of Gaussians with diagonal covariance, Jiang et al. (2017) gave a closed-form solution to compute such an expectation (see Corollary 5.1 in Supplementary Material 5). For general Gaussian distributions, we give the closed-form solution in Lemma 5.2 of Supplementary Material 5, but, for simplicity, we assume Gaussian distributions with diagonal covariances throughout this work.

Our algorithmic framework thus consists of (a) a multimodal VAE with a GMM prior and (b) a parameterization of our DAG and BN that ties into the GMM of (a), as highlighted in Figure 1. We describe each of these components in the subsections below.

2.1.1 Multimodal VAE with a GMM prior

Our multimodal representation learning framework amounts to a VAE with a GMM prior on the latent space; see the orange box (a) in Figure 1. Specifically, the latent space is assumed to be a mixture of Gaussian distributions, where is the number of nodes and is the number of outcomes of the node . The joint probability on the mixture assignment in the GMM is additionally assumed to factorize with a Bayesian network prior to obey the Markov factorization property in Equation 1. Assignment to an individual cluster has the probabilityand each cluster in the GMM has a Gaussian distributionwhere the parameters and are recovered by training with a first-order optimizer or computed using a block-coordinate maximization strategy outlined in Section 2.2.2.

The multimodal embedding and decoding of the VAE is similar to Walker et al. (2024). In particular, we use neural network encoders to embed each modality as a Gaussian and then combine these embeddings using a product of experts (see in Figure 1). In other words, for each modality , we assume , where for a neural network with trainable parameters . We deterministically compute the multimodal embedding from the unimodal ones via the identitywhere is a normalization constant, andDuring training, the multimodal distribution is sampled using the reparameterization trick. In other words, we sample and compute , where is the Hadamard product.

Our decoders output a Gaussian for each modality . The Gaussians’ parameters are determined by neural networks , i.e., . Alternatively our decoders can be expert models, or they can depend upon only , i.e., .

2.1.2 Directed acyclic graph and joint distribution of nodes

The GMM prior in the previous section implicitly uses the Markov factorization property, which requires knowledge of a DAG to relate the causal dependencies between clusters; since we do not assume knowledge of this DAG a priori, we must recover this DAG and the accompanying probability distribution of the resulting BN while simultaneously training our VAE. We first describe our DAG parameterization, followed by our method to compute the resulting joint probability distribution.

Our DAG parameterization builds off of concepts from Hodge theory (Jiang et al., 2011; Lim, 2020). We parameterize an edge indicator function as the graph gradient on a set of nodes. By using the graph gradient, we guarantee that our edge indicator function is curl-free, and consequently, we define a complete DAG; we introduce sparsity in the complete DAG through nonnegative weightings of edges.

As a brief summary, Hodge theory provides a generalization for vector calculus concepts of gradient, curl, and divergence, and discrete Hodge theory generalizes these notions to the discrete setting of graphs to define the graph gradient, graph curl, graph divergence, etc. Through these definitions, we can view numeric labels on nodes and edges as functions on the nodes and edges, respectively. If we define the functions on the edges to be the graph gradient of the edge’s adjacent node values, then we view the edges as the flow between the nodal values. Furthermore, by defining edge values via the graph gradient, we ensure that our edges define a curl-free function, which means that our edge values define a DAG since curl-free functions on graphs have no cycles by construction. We note that the graph gradient will define a complete DAG, so we use trainable nonnegative coefficients to introduce sparsity. For more information on Hodge theory and the generalization to graphs, see Trask et al. (2020); Lim (2020).

Explicitly, given a set of nodes, our DAG parameterization assigns a trainable scalar score to each node . We denote the vector of trainable node scores as . Each potential edge between nodes is assigned a value given bywhere is the graph gradient operator, and each is a trainable nonnegative scalar used to induce sparsity in . In terms of Hodge theory, one can view as an trainable nonnegative metric tensor and as a flow from node to node .

We use these edge values to give a regularized edge indicator function:where the scalar is a temperature parameter that controls the sharpness of the regularization of the indicator function. We use this temperature to control how easily the DAG can update during training, as described in Section 2.2.3. With this formulation, we assign edges in our DAG via the rule

Theorem 2.1. Let A be the adjacency matrix of a directed graph G. Then G is a DAG if and only if for some matrix E with entries given by Equation 3.

Our DAG parameterization does guarantee a DAG, and parameterization is flexible enough to learn any possible DAG. Formal proofs are given in Supplementary Material 3.

Proof: see Lemma 3.4 and Proposition 3.6 in Supplementary Material 3.

We now proceed to compute the joint probability distribution on as per the Markov factorization property. This means that we must parametrize the probability distribution at each node and have a method for computing from . Our parametrization must be flexible enough to allow for the direction of the edge dependencies in the DAG to change during training, and, as a result, the number of causal factors that are parents of a given node (i.e., ) may change as well. We therefore build, for each node , a parameterization of these probabilities that allows for any subset of nodes to be parents via a trainable tensor ; we downselect which nodes are parents by using the DAG edge indicator scores from Equation 3 and averaging out those modes that are not parents. For ease of notation, we let .

For each , we define to be a nonnegative rank- tensor of size constrained so that for any ,This constraint ensures that the entries of represent the probabilitieswhere is a one-hot encoding, with the entry set to 1.

With this definition, we now proceed to describe our downselection algorithm for parameterizing the probabilities over the structure of a given DAG. If a node is not a parent node of , then we remove mode from by contracting against (where is the vector of all ones) along mode . In essence, we are replacing the mode in by its average. This contraction makes independent of when is not a parent of . If the node is a parent of , then we can contract against the realizations of the parent one-hot encodings since they are already known. We, therefore, end up with an expression for the categorical distribution on viawhere denotes the contractive -mode tensor product against mode (see Kolda and Bader (2009); Bader and Kolda (2006) for more details), andUsing our DAG representation, these cases can be written as follows:or more concisely aswhich has the benefit of allowing us to handle relaxations of when is not necessarily a binary matrix (such as when the temperature is small but not yet sufficiently close to 0). With defined as such, we can write via

From the tensors and the vectors , we can now compute .

By Corollary 3.5, we may assume that the categorical variables are ordered such that , ; if not, we reassign the indices via the permutation . Observe that, by assumption of the Markov factorization property,where and where the expression for is given by in Equation 5. This inductive process of computing is given in Algorithm 1.

As a result of Algorithm 1, we have flexible parameterization for the joint distribution on , which is used for the GMM mixing probabilities in our VAE. This parameterization, of both and the underlying DAG, are differentiable and amenable to automatic differentiation, allowing us to train in an end-to-end fashion. Our training approach is described in the next section.

2.2 ELBO loss and training

We now describe the loss function used to train pimaDAG. The ELBO loss for training is similar to that of Walker et al. (2024), albeit with a different notation and an additional computation of cluster assignment based on the DAG.

2.2.1 Single-sample ELBO

The full ELBO derivation is in Supplementary Material 4. After dropping constant terms, the single-sample ELBO is

where is an estimate for the posterior distribution and, following Jiang et al. (2017), is computed bywhere we recall for convenience. Note that and are both tensors with modes of size . The tensor can be calculated via Algorithm 1 in Section 2.1.2. The values of can be computed by sampling from each Gaussian in the Gaussian mixture model. All other values are parameters in our model. Table 2 summarizes the assumed distributions on each term in the architecture and lists how the variables are computed and updated during training.

2.2.2 Training

To train our causal model, we seek to maximize the ELBO over the entire dataset. In other words, if we use to denote Equation 6 for the datapoint, then we want to minimize . Throughout training, we alternate between (1) updating the neural network, expert model, and DAG parameters via gradient descent and (2) updating the Gaussian mixture centers and variances using block-coordinate maximization, similar to Walker et al. (2024). In particular, we compute the optimal Gaussian mixture centers and variances by taking the derivative of with respect to these parameters and solving for the global minimizers:where indexes the data point, and, in particular, and are, respectively, the encoded mean and variance of the data point. Our training procedure follows Algorithm 2.

2.2.3 Practical considerations

We implement several tools for aiding training and algorithm adaptation. These tools are described here, and the use of these tools in each experiment is detailed in Supplementary Material 2.

2.2.3.1 Pre-training

We sometimes found that before fitting a DAG, we need a reasonable latent embedding, in the sense that the DAG is meant to relate informative latent features instead of random features at initialization. Thus, we implemented an optional pre-training regimen, following Jiang et al. (2017) and Kingma and Welling (2014). As our first step in pre-training, we fix the pre-initialized encoders and initialize the cluster means and variances via Equations 7, 8. This initialization has the benefits of providing a good GMM fit for the initial latent embedding, but if the initial latent embedding is poor or undiscriminating, then the initial GMM fitting by these steps might not focus on any informative features. Our next step in pre-training is to train the weights and biases of the encoders and decoders via the reconstruction term or by fitting a unit-normal Gaussian variational autoencoder (Kingma and Welling, 2014). Following this training, we find a good initial GMM fit through iterations of Equations 7, 8. This has the benefit of finding a good initial embedding, from which block-coordinate maximization can recover meaningful features.

2.2.3.2 Edge indicator function adaptations

We have two optional adaptations to the edge indicator function. The first is to add random noise to the node scores . This noise is included to break free of local minima and may additionally test edge orientation. Our second adaptation is to anneal during training. In Equation 3, serves as a temperature parameter, and, as , approaches a true indicator function. Our annealing implementation is simple, where we specify the initial , the final , and the update frequency of .

2.2.3.3 Updates on GMM parameters

The cluster center and variance updates in Equation 8, paired with the gamma calculation in Equation 7, are reminiscent of expectation-maximization. The traditional maximization step would, however, also update the probability of belonging to each cluster . Although there is no closed-form expression for an update on from our ELBO since they depend on the underlying causal factorization, we alternatively perform extra gradient-descent steps to update after each update of the cluster means and variances. Furthermore, we also implemented the option to perform multiple iterations of GMM variable updates per epoch.

3 Results

Although we prove in Proposition 3.6 that our DAG parameterization is capable of learning any DAG, we demonstrate this capability in Section 3.1. We then test pimaDAG on a synthetic dataset consisting of circle images and a scientific dataset consisting of 3D-printed lattices. All architectures, hyperparameters, and training details for the experiments can be found in Supplementary Material 2.

3.1 Efficacy of DAG parameterization

We test the ability of our DAG parameterization (Equation 3) to recover any DAG. First, we generate a random set of DAGs of various sizes, i.e., the number of nodes , by generating a (random) consistent ordering, inducing a DAG on the complete graph of nodes, and then randomly discarding edges from the complete graph by a Bernoulli distribution with probability that an edge is discarded. For each graph in our dataset, we train to recover the specified DAG’s adjacency matrix, comparing the edge adjacency matrix of the predicted graph to that of the true graph using binary cross-entropy (BCE) as our loss and Adam as our optimizer. We use a learning rate of and train for steps with a batch size of 16. For the largest size , training was extended to steps as the loss had not converged after the first steps; this slowdown in convergence is expected, given that the number of unique DAGs grows superexponentially with the number of vertices in a graph (Kuipers and Moffa, 2015). Training begins with an initial temperature of 0.5, and we reduce by 1% every 100 steps to a minimum temperature of . We repeat this test for for each size graph. In Table 3, we report the average and standard deviation of the BCE loss for each of the sampled DAGs for each size. In every sample at every size in our test, with the exception of a single sample for , the largest graph size, we recover the adjacency matrix within entry-wise accuracy, often with accuracies of or smaller; the maximum absolute entry-wise error (MAE) in the predicted adjacency graph across all entries and all samples of that size is also reported in Table 3. Similar statistics are shown for common test DAGs, taken from the BNLearn (Scutari, 2010) repository of Bayesian networks, yielding comparable results, as shown in Table 4.

We additionally compare our ability to recover the DAG adjacency matrix against NO TEARS (Zheng et al., 2018) on the same benchmarks in Table 4 using the same metrics for success used in Table 3. To make pimaDAG and NO TEARS comparable, we do not perform any optional enhancement of the discovered DAG in Zheng et al. (2018), such as weight thresholding or regularization, and we otherwise use the same parameterization, as provided for in their original work. We note a handful of key differences between these two methods, namely, that (1) pimaDAG’s parameterization yields an unconstrained nonlinear optimization problem, while NO TEARS yields a linear problem (depending on choice of loss, as described in Zheng et al. (2018)) with a nonlinear constraint; (2) pimaDAG uses a first-order optimizer, while NO TEARS uses a second-order optimizer as part of an augmented Lagrangian scheme; and (3) pimaDAG always guarantees a DAG by construction, while NO TEARS guarantees a DAG only when the constraints are satisfied, which is not guaranteed a priori in the augmented Lagrangian formulation. Thus, we see in Table 4 that when NO TEARS successfully finds a DAG, it generally does so more efficiently and accurately, which is unsurprising given its more sophisticated optimization approach but that for the larger graphs tested from the BNLearn repository, NO TEARS fails to recover a DAG and generally performs poorly on these problems. This behavior motivates the need for methods like pimaDAG, which guarantee a DAG by construction during all stages of end-to-end training.

3.2 Circles

For our first pimaDAG experiment, we generated a synthetic unimodal dataset consisting of images of circles with three different features: hue (red and blue), radius (Gaussian mixture of big and small), and shift (Gaussian mixture of left and right). We generated 4,096 circles using the decision tree in Supplementary Figure S2 where we purposefully overlapped distributions of , , and to necessitate the discovery of a DAG describing the generative process. We ran this experiment with three nodes in the DAG, where each node was a binary categorical random variable. The latent space showed disentanglement in the three different features. The learned DAG is in subpanel (E) of Figure 2. By comparing cluster labels to features characteristic of each cluster, we see that node in the DAG corresponds to the radius, node corresponds to the hue, and node corresponds to the shift. For example, we conclude that corresponds to the hue as all clusters with red circles have a 0 in the second entry of their label. Under this identification, the resulting directed acyclic graph demonstrates that the radius and hue play a key role in the outcome of the shift.

FIGURE 2

3.3 Lattices

Our next experiment uses a dataset of 3D-printed lattices (Garland et al., 2020). In this dataset, two different lattice geometries were printed (octet and gyroid) from 316 L stainless steel, where each respective geometry was printed using the same 3D model but different print process parameters. Each lattice was a 10-mm cube with three unit cells. Strut diameters for the octets and wall thicknesses for the gyroids were both 0.5 mm in the model design. Although wall thicknesses and strut diameters were set to 0.5 mm, actual thicknesses and diameters varied due to the varied print process parameters. This dataset included a total of 91 printed lattices, where two modalities for each lattice sample were collected: an image of the lattice and a stress–strain curve produced by a high-throughput uniaxial compression machine . The stress–strain curves represent a physics-imbued modality, where curves can be modeled via continuous piecewise linear functions; see Trask et al. (2022). Consequently, for the stress–strain modality, we used an expert model decoder composed of two piecewise linear segments. Specifying two binary feature nodes resulted in a latent space organized by the lattice type and stress–strain curves, see Figure 3. The two clusters consisting of the octet geometry merged, and the corresponding expert models are nearly identical. This result is consistent with the distribution of octet stress–strain curves, which has a lower variance than the gyroid stress–strain curves. By comparing cluster labels to features characteristic of each cluster, we see that node corresponds to the lattice type while corresponds to the stress–strain curve profile. The learned DAG suggests that the lattice type influences the stress–strain curve.

FIGURE 3

By way of comparison, note that Trask et al. (2022) tested PIMA on this same lattice dataset. In their Figure 6, they found two clusters: one for each lattice type. Their latent space provided a disentanglement of the data but did not give the additional insight given by our DAG, which identifies dependencies between stress–strain profiles and the lattice type. We thus anticipate that pimaDAG can aid in the discovery of hidden, dependent features in scientific datasets.

4 Conclusion

We present a general-purpose framework, pimaDAG, for discovering a Bayesian network of latent features in high-dimensional data. This framework is capable of handling multiple modalities and physics constraints to encourage disentanglement of data with a conditional independence structure. We introduce a new parameterization for learning DAGs, and we prove that this parameterization is capable of discovering any DAG; for a selection of DAGs, we confirm this via numerical experiments. We demonstrate the efficacy of pimaDAG on synthetic and real data, and we were able to achieve interpretable dependent features. These results show that meaningful disentanglement via a trainable Bayesian network is possible, even in purely exploratory settings.

There are some limitations to this framework. First, our model assumes discrete features, which means it may be less suitable for continuous features. Next, variational autoencoders are notoriously difficult to train to identify features amenable to analysis or that match extant intuition. In part, this follows from the fact that variational autoencoders only provide interpretable latent representations under certain constraints. Although the aim of this work was to discover a Bayesian network and its feature nodes for general data, with perhaps multiple modalities and physics constraints to provide pseudo-labels, future work will investigate the conditions necessary to generate unique features and conditional relationships within pimaDAG. We furthermore plan to investigate the possibility of including interventional data.

References

• 1

  BaderB. W.KoldaT. G. (2006). Algorithm 862: matlab tensor classes for fast algorithm prototyping. ACM Trans. Math. Softw. (TOMS)32, 635–653. 10.1145/1186785.1186794

  • CrossRef
  • Google Scholar
  • View reference in article

• 2

  BengioY.CourvilleA.VincentP. (2013). Representation learning: a review and new perspectives. IEEE Trans. Pattern Analysis Mach. Intell.35, 1798–1828. 10.1109/TPAMI.2013.50

  • CrossRef
  • Google Scholar
  • View reference in article

• 3

  BoyceB. L.UchicM. D. (2019). Progress toward autonomous experimental systems for alloy development. MRS Bull.44, 273–280. 10.1557/mrs.2019.75

  • CrossRef
  • Google Scholar
  • View reference in article

• 4

  BuchholzS.RajendranG.RosenfeldE.AragamB.Sch olkopfB.RavikumarP. (2023). Learning linear causal representations from interventions under general nonlinear mixing. NeurIPS.

  • Google Scholar
  • View reference in article

• 5

  ChickeringD. M.HeckermanD.MeekC. (2004). Large-sample learning of bayesian networks is np-hard. J. Mach. Learn. Res.

  • Google Scholar
  • View reference in article

• 6

  DilokthanakulN.MedianoP. A.GarneloM.LeeM. C.SalimbeniH.ArulkumaranK.et al (2016). Deep unsupervised clustering with Gaussian mixture variational autoencoders. arXiv Prepr. arXiv:1611.02648. 10.48550/arXiv.1611.02648

  • CrossRef
  • Google Scholar
  • View reference in article

• 7

  GaoT.BhattacharjyaD.NelsonE.LiuM.YuY. (2022). Idyno: learning nonparametric dags from interventional dynamic data. Int. Conf. Mach. Learn.

  • Google Scholar
  • View reference in article

• 8

  GarlandA. P.WhiteB. C.JaredB. H.HeidenM.DonahueE.BoyceB. L. (2020). Deep convolutional neural networks as a rapid screening tool for complex additively manufactured structures. Addit. Manuf.35, 101217. 10.1016/j.addma.2020.101217

  • CrossRef
  • Google Scholar
  • View reference in article

• 9

  GlymourC.ZhangK.SpirtesP. (2019). Review of causal discovery methods based on graphical models. Front. Genet.10, 524. 10.3389/fgene.2019.00524

  • CrossRef
  • Google Scholar
  • View reference in article

• 10

  HigginsI.MattheyL.PalA.BurgessC.GlorotX.BotvinickM.et al (2017). “beta-vae: learning basic visual concepts with a constrained variational framework,” in 5th international conference on learning representations, ICLR 2017.

  • Google Scholar
  • View reference in article

• 11

  HyvarinenA.MoriokaH. (2017). “Nonlinear ica of temporally dependent stationary sources,” in Artificial intelligence and statistics (PMLR), 460–469.

  • Google Scholar
  • View reference in article

• 12

  JensenF. V. (2001). Bayesian networks and decision graphs. Springer.

  • Google Scholar
  • View reference in article

• 13

  JiangX.LimL.-H.YaoY.YeY. (2011). Statistical ranking and combinatorial hodge theory. Math. Program.127, 203–244. 10.1007/s10107-010-0419-x

  • CrossRef
  • Google Scholar
  • View reference in article

• 14

  JiangZ.ZhengY.TanH.TangB.ZhouH. (2017). “Variational deep embedding: an unsupervised and generative approach to clustering,” in Proceedings of the 26th international joint conference on artificial intelligence, 1965–1972.

  • Google Scholar

• 15

  KalainathanD.GoudetO.GuyonI.Lopez-PazD.SebagM. (2022). Structural agnostic modeling: adversarial learning of causal graphs

  • Google Scholar

• 16

  KhattarD.GoudJ. S.GuptaM.VarmaV. (2019). “Mvae: multimodal variational autoencoder for fake news detection,” in The world wide web conference, 2915–2921.

  • Google Scholar

• 17

  KingmaD. P.WellingM. (2014). “Auto-encoding variational bayes,” in 2nd international conference on learning representations, ICLR 2014.

  • Google Scholar
  • View reference in article

• 18

  KingmaD. P.WellingM. (2019). An introduction to variational autoencoders. Found. Trends® Mach. Learn.12, 307–392. 10.1561/2200000056

  • CrossRef
  • Google Scholar
  • View reference in article

• 19

  KocaogluM.SnyderC.DimakisA. G.VishwanathS. (2018). “CausalGAN: learning causal implicit generative models with adversarial training,” in 6th international conference on learning representations, ICLR.

  • Google Scholar
  • View reference in article

• 20

  KoldaT. G.BaderB. W. (2009). Tensor decompositions and applications. SIAM Rev.51, 455–500. 10.1137/07070111x

  • CrossRef
  • Google Scholar
  • View reference in article

• 21

  KuipersJ.MoffaG. (2015). Uniform random generation of large acyclic digraphs. Statistics Comput.25, 227–242. 10.1007/s11222-013-9428-y

  • CrossRef
  • Google Scholar
  • View reference in article

• 22

  LeeH.-C.DanielettoM.MiottoR.CherngS. T.DudleyJ. T. (2020). Scaling structural learning with no-bears to infer causal transcriptome networks. Biocomputing25, 391–402. 10.1142/9789811215636_0035

  • CrossRef
  • Google Scholar
  • View reference in article

• 23

  LimL.-H. (2020). Hodge laplacians on graphs. Siam Rev.62, 685–715. 10.1137/18m1223101

  • CrossRef
  • Google Scholar
  • View reference in article

• 24

  LippeP.MagliacaneS.L oweS.AsanoY. M.CohenT.GavvesE. (2022). “Citris: causal identifiability from temporal intervened sequences,” in 39th international conference on machine learning.

  • Google Scholar
  • View reference in article

• 25

  LippeP.MagliacaneS.L oweS.AsanoY. M.CohenT.GavvesE. (2023a). “Biscuit: causal representation learning from binary interactions,” in Thirty-ninth conference on uncertainty in artificial intelligence.

  • Google Scholar
  • View reference in article

• 26

  LippeP.MagliacaneS.L oweS.AsanoY. M.CohenT.GavvesE. (2023b). “Causal representation learning for instantaneous and temporal effects in interactive systems,” in International conference on learning representations.

  • Google Scholar
  • View reference in article

• 27

  LöweS.MadrasD.ZemelR.WellingM. (2022). “Amortized causal discovery: learning to infer causal graphs from time-series data,” in Conference on causal learning and reasoning (PMLR), 509–525.

  • Google Scholar
  • View reference in article

• 28

  LyuQ.FuX.WangW.LuS. (2022). Understanding latent correlation-based multiview learning and self-supervision: an identifiability perspective. ICLR.

  • Google Scholar
  • View reference in article

• 29

  MoriokaH.HyvarinenA. (2023). “Connectivity-contrastive learning: combining causal discovery and representation learning for multimodal data,” in International conference on artificial intelligence and statistics (PMLR), 3399–3426.

  • Google Scholar
  • View reference in article

• 30

  NgI.ZhuS.ChenZ.FangZ. (2019). A graph autoencoder approach to causal structure learning. NeurIPS 2019 Workshop.

  • Google Scholar
  • View reference in article

• 31

  PamfilR.SriwattanaworachaiN.DesaiS.PilgerstorferP.BeaumontP.GeorgatzisK.et al (2020) “Dynotears: structure learning from time-series data,” in 23rd international conference on artificial intelligence and statistics.

  • Google Scholar
  • View reference in article

• 32

  RamseyJ.GlymourM.Sanchez-RomeroR.GlymourC. (2017). A million variables and more: the fast greedy equivalence search algorithm for learning high-dimensional graphical causal models, with an application to functional magnetic resonance images. Int. J. Data Sci. Anal.3, 121–129. 10.1007/s41060-016-0032-z

  • CrossRef
  • Google Scholar
  • View reference in article

• 33

  SchölkopfB.LocatelloF.BauerS.KeN. R.KalchbrennerN.GoyalA.et al (2021). Toward causal representation learning. Proc. IEEE109, 612–634. 10.1109/jproc.2021.3058954

  • CrossRef
  • Google Scholar
  • View reference in article

• 34

  ScutariM. (2010). Learning bayesian networks with the bnlearn r package. J. Stat. Softw.35. 10.18637/jss.v035.i03

  • CrossRef
  • Google Scholar
  • View reference in article

• 35

  ShenX.LiuF.DongH.LianQ.ChenZ.ZhangT. (2022). Weakly supervised disentangled generative causal representation learning

  • Google Scholar
  • View reference in article

• 36

  ShimizuS.HoyerP. O.HyvärinenA.KerminenA. (2006). A linear non-Gaussian acyclic model for causal discovery. J. Mach. Learn. Res.7, 2003–2030.

  • Google Scholar
  • View reference in article

• 37

  SparkesA.AubreyW.ByrneE.ClareA.KhanM. N.LiakataM.et al (2010). Towards robot scientists for autonomous scientific discovery. Autom. Exp.2, 1. 10.1186/1759-4499-2-1

  • CrossRef
  • Google Scholar
  • View reference in article

• 38

  SpirtesP.GlymourC.ScheinesR. (2000) Causation, prediction, and search, 81. The MIT Press.

  • Google Scholar

• 39

  SquiresC.SeigalA.BhateS.UhlerC. (2023). “Linear causal disentanglement via interventions,” in International conference on machine learning (PMLR).

  • Google Scholar
  • View reference in article

• 40

  TraskN.HuangA.HuX. (2020). Enforcing exact physics in scientific machine learning: a data-driven exterior calculus on graphs. arXiv preprint arXiv:2012.11799

  • Google Scholar
  • View reference in article

• 41

  TraskN.MartinezC.LeeK.BoyceB. (2022). Unsupervised physics-informed disentanglement of multimodal data for high-throughput scientific discovery. arXiv Prepr. arXiv:2202.03242. 10.48550/arXiv.2202.03242

  • CrossRef
  • Google Scholar
  • View reference in article

• 42

  VariciB.AcartürkE.ShanmugamK.KumarA.TajerA. (2023). “Score-based causal representation learning with interventions,” in Causal representation learning workshop at NeurIPS.

  • Google Scholar
  • View reference in article

• 43

  VowelsM.CamgozN. C.BowdenR. (2022). D’ya like dags? a survey on structure learning and causal discovery. ACM Comput. Surv.55, 1–36. 10.1145/3527154

  • CrossRef
  • Google Scholar
  • View reference in article

• 44

  WalkerE.TraskN.MartinezC.LeeK.ActorJ. A.SahaS.et al (2024). Unsupervised physics-informed disentanglement of multimodal data. Found. Data Sci.0, 0. 10.3934/fods.2024019

  • CrossRef
  • Google Scholar

• 45

  WangL.HuangS.WangS.LiaoJ.LiT.LiuL. (2024). A survey of causal discovery based on functional causal model. Eng. Appl. Artif. Intell.133, 108258. 10.1016/j.engappai.2024.108258

  • CrossRef
  • Google Scholar
  • View reference in article

• 46

  WeiD.GaoT.YuY. (2020). “Dags with no fears: a closer look at continuous optimization for learning bayesian networks,” in Conference on neural information processing systems.

  • Google Scholar
  • View reference in article

• 47

  YangM.LiuF.ChenZ.ShenX.HaoJ.WangJ. (2021). “Causalvae: structured causal disentanglement in variational autoencoder,” in Conference on computer Vision and pattern recognition (IEEE/CVF).

  • Google Scholar
  • View reference in article

• 48

  YaoD.XuD.LachapelleS.MagliacaneS.TaslakianP.MartiusG.et al (2024). Multi-view causal representation learning with partial observability. ICLR.

  • Google Scholar
  • View reference in article

• 49

  YaoW.ChenG.ZhangK. (2022). Temporally disentangled representation learning. Adv. Neural Inf. Process. Syst.35, 26492–26503.

  • Google Scholar

• 50

  YuY.ChenJ.GaoT.YuM. (2019). “Dag-gnn: dag structure learning with graph neural networks,” in International Conference on machine learning (PMLR).

  • Google Scholar

• 51

  YuY.GaoT.YinN.JiQ. (2021). “Dags with no curl: an efficient dag structure learning approach,” in International conference on machine learning (PMLR), 12156–12166.

  • Google Scholar

• 52

  ZhengX.AragamB.RavikumarP.XingE. P. (2018). Dags with no tears: continuous optimization for structure learning. Adv. Neural Inf. Process. Syst.31.

  • Google Scholar