Primary Submission Category: Machine Learning and Causal Inference
A new parametrization of DAGs and causal Markov kernels for scientific feature discovery
Authors: Elise Walker, Jonas Actor, Carianne Martinez, Nathaniel Trask,
Presenting Author: Elise Walker*
Due to the complexity of multimodal scientific datasets, causal feature detection of these datasets necessitates unsupervised representation learning methods. Physics-informed multimodal autoencoders (PIMA) have recently demonstrated successful unsupervised feature detection with variational autoencoders (VAEs) where multiple scientific modalities and physics constraints acted as surrogates for the supervision labels typically needed for successful VAEs. Building upon the successes of PIMA, we present a VAE framework coupled with a trainable directed acyclic graph (DAG) to discover features with plausible causal relationships in multimodal scientific datasets. In particular, we introduce a new parametrization for learning both the edges of a DAG and the causal Markov kernels of the joint distribution of its nodes. We use this parametrization to simultaneously learn a DAG in conjunction with a latent space of a VAE. Training of our DAG and VAE is performed in an end-to-end differentiable framework via a single, tractable evidence lower bound (ELBO) loss function. We achieve a single ELBO by placing a Gaussian mixture prior on the latent space and identifying each of the Gaussians with an outcome of the DAG nodes. We demonstrate the efficacy of our DAG parametrization, and we test our joint VAE and DAG framework on both a synthetic and a scientific dataset. Our results demonstrate the capability of learning a DAG on discovered key features in an exploratory scientific setting.