Spatial Transcriptomics Learned Non-negative Matrix Factorization
STL NMF is the culmination of my research in Non-negative Matrix Factorization (NMF) and its application to Spatial Transcriptomics data. This work took place at Case Western Reserve University under the guidance of Dr. Weihong Guo, with the endless support of fellow researcher Jiasen Zhang. Read more about the method and resulting paper below.
Background
Spatial transcriptomics is the study of spatial transcriptomics data (X) — a count matrix of m genes across n spatial locations. This ST data is often aided by an RNA-seq reference dataset (B) — an expression matrix of the same m genes across p single cells.
Biologically, genes and cells are linked through both position and function. Using this, Patrik Ståhl et al. first proposed using Spatial Transcriptomics data to deconvolve a spatial mapping of genes into a spatial mapping of cell types (V).
There are many strategies to do this, each with strengths and weaknesses — the most popular being deep-learning, probabilistic, optimal-transport, and NMF-based methods. My research focuses on NMF.
Deconvolution
To set up a matrix-factorization solution, we start from the equation X = BV + ε. From here we aim to find the "true" V, given both X and B. We then use the fact that all three matrices are non-negative (every entry is zero or a positive real), which gives us a way to solve for V (eq. 2).
However, this leaves an ill-posed optimization problem that struggles to find unique solutions. Luckily, we can apply biological priors as regularization terms to guide it toward an optimal solution.
Locations that are close together and within the same biological region should have far more similar genetic makeups than distant spots. We encode this with an adjacency term in the objective that limits NMF's freedom, enforcing similarity across nearby locations and disparity across distant ones. We also know the total proportion of all cell types in one location should sum to one — a sum-to-one constraint on the columns of V.
After adding these terms (plus one more, mostly inconsequential), we arrive at the objective (eq. 3) and its multiplicative update rule (eq. 4).
With the problem defined, we can solve it: start from an initialized (random) matrix V and iteratively apply the update rule (eq. 4). Over time V converges to a stable solution that best approximates the "true" V. Even so, this solution is imperfect — it relies heavily on good choices of hyperparameters (λ₁, λ₂, λ₃, W).
Algorithm Unrolling
To address that reliance, we use a strategy called algorithm unrolling, first proposed by Gregor and LeCun in 2010. It transforms our iterative optimization into a (repetitive) single-layer neural network that trains on the problem's hyperparameters. Doing so, we learn the optimal hyperparameters (λ₁, λ₂, λ₃, W) and push our solution past many state-of-the-art methods.