Nearest-neighbor nonnegative spatial factorization to study spatial and temporal transcriptomics

Nonnegative spatial factorization (NSF) is a spatially-aware factorization method that uses Gaussian processes (GPs) as spatial priors in a Poisson latent factor model to robustly identify interpretable, parts-based representations in spatial transcriptomics data. However, NSF scales poorly with modern datasets due to the computational complexity of Gaussian processes, which scales cubically with the number of points used for inference O(N 3). To address this limitation, we propose a modified version of NSF that leverages variational nearest neighbor Gaussian processes (VNNGPs), resulting in a substantial reduction in inference complexity from O(NM 2) in the current version of NSF to O(MNK2) for M inducing points, N total points and K nearest neighbors. Our method, nearest-neighbor NSF (NNNSF), is benchmarked on synthetic and real-world spatial and temporal transcriptomics datasets. Experimental results demonstrate that NNNSF achieves linear scaling with the number of neighbors and points used for inference in contrast with NSF, which has exponential computational complexity as the number of points used in inference increases. By restricting covariance calculations to the K-nearest neighbors of the points used in inference, NNNSF allows the use of more inducing points, leading to lower reconstruction loss. Nearest-neighbor NSF (NNNSF), which replaces standard variational inference with inducing points in NSF with the VNNGP, leads to a computationally efficient and scalable version of NSF that can be applied to large existing and forthcoming spatial genomics data. We added VNNGP and NNNSF to the GPZoo package, an an ongoing open source project developing a modular Gaussian process library in Python making use of the PyTorch interface. Source code and demonstrations are available at https://github.com/luisdiaz1997/GPzoo/tree/main.