Making Multi-Axis Gaussian Graphical Models Scalable to Millions of Samples and Features
Bailey Andrew, David R. Westhead, Luisa Cutillo
TL;DR
This work tackles scalable inference of conditional dependencies in data that do not satisfy sample independence, focusing on multi-axis, tensor-variate datasets. It develops a scalable, independence-free Gaussian graphical model based on a singular Kronecker-sum normal distribution with low-rank axis graphs, achieving $O(n^2)$ time and $O(n)$ space. The approach preserves multi-modality and arbitrary marginals via a Gaussian copula, provides interpretable hyperparameters, and enables edge-wise hypothesis testing using the Fisher information. Demonstrations on synthetic data and large real-world datasets, including a million-cell scRNA-seq PBMC dataset, illustrate significant scalability gains and competitive accuracy relative to prior methods.
Abstract
Gaussian graphical models can be used to extract conditional dependencies between the features of the dataset. This is often done by making an independence assumption about the samples, but this assumption is rarely satisfied in reality. However, state-of-the-art approaches that avoid this assumption are not scalable, with $O(n^3)$ runtime and $O(n^2)$ space complexity. In this paper, we introduce a method that has $O(n^2)$ runtime and $O(n)$ space complexity, without assuming independence. We validate our model on both synthetic and real-world datasets, showing that our method's accuracy is comparable to that of prior work We demonstrate that our approach can be used on unprecedentedly large datasets, such as a real-world 1,000,000-cell scRNA-seq dataset; this was impossible with previous approaches. Our method maintains the flexibility of prior work, such as the ability to handle multi-modal tensor-variate datasets and the ability to work with data of arbitrary marginal distributions. An additional advantage of our method is that, unlike prior work, our hyperparameters are easily interpretable.