Praxis-BGM: Prior-Augmented and Regularized Natural-Gradient Variational Inference with Gaussian Mixture Models for Transfer Learning

Music City Center 

Recent advances in high-throughput omics technologies have enabled observational studies to collect multiple omics layers on the same individuals. However, high-dimension data and often low-sample-sizes (HDLSS) pose significant challenges for model-based clustering approaches such as Gaussian Mixture Models (GMMs). While existing methods for integrative multi-omics clustering account for variations within and across omic layers, they often fall short in addressing the HDLSS issue, where complex mixture patterns become difficult to generalize in a small number of subjects and model instability increases as model complexity increases.
Statistical transfer learning has emerged as a powerful approach to address HDLSS by leveraging knowledge from related but distinct source domains to improve modeling in the target domain. Among various strategies, incorporating informative priors from the source domain within a Bayesian framework offers a natural and effective solution. In addition, modern Bayesian methods offer scalable and efficient computation for high-dimensional data. For example, natural-gradient variational inference turns Bayesian inference into an optimization problem and leverages the underlying geometry of the parameter space to achieve fast yet good posterior approximation.
We introduce Praxis-BGM, a natural-gradient variational inference method for GMMs that flexibly incorporates cluster-specific priors—including means, covariance matrices, and structural pathway information—to constrain posterior estimation and facilitate knowledge transfer. These various prior components can be used individually or in combination depending on their availability. They can be obtained from large source datasets or reference atlases. For estimation, we optimize the variational covariance matrices in the Cholesky decomposition space, which ensures positive definiteness, enhances numerical stability, and reduces the number of free parameters for efficiency. We derive natural-gradient updates that incorporate prior knowledge and propose a clustering-driven feature selection procedure based on Bayes Factors. Praxis-BGM is implemented using the JAX library for accelerator-oriented computation with high efficiency and scalability.
We demonstrate the effectiveness of Praxis-BGM through extensive simulations, evaluating the contribution of each component of the informative prior, and assessing the method's ability to overcome less accurate priors while balancing inconsistencies between the priors and observed data. We also demonstrate the application of Praxis-BGM to three applied analyses: 1) two similar COVID-19 metabolomics datasets; 2) two breast cancer transcriptomic datasets from The Cancer Genome Atlas (TCGA) and the Molecular Taxonomy of Breast Cancer International Consortium (METABRIC); and 3) a single-cell RNA sequencing data set using atlas reference data as the prior. The first two applications highlight how priors for meaningful cluster structures derived from one source study can enhance clustering performance in another, improving both biological interpretability and clinical relevance. In the third example, mean priors derived from labeled cell types in a reference atlas are used to predefine and annotate clusters in the observed data, guiding the estimation process.

Multi-omics

Bayesian Clustering

Mixture Model

Variational Inference

High-Dimensional Data

Dimension Reduction 

Section on Statistical Learning and Data Science