Find your microenvironments faster with Neural Spatial LDA
Abstract
Spatial organization of different cell types in tissues have been shown to be important factors in many important biological processes such as aging, infection and cancer [3].
In particular, organization of the cells in a tumor microenvironment (TME) has been shown to play a crucial role in treatment response, disease pathology and patient outcome [14].
Spatial LDA [5] is a general purpose probabilistic model that has been used to discover novel microenvironments in settings such as Triple Negative Breast Cancer (TNBC) and Tuberculosis infections.
However, the implementation of Spatial LDA proposed in [5] uses variational inference for learning model parameters and unfortunately does not scale well with dataset size and does not lend itself to speed-up via GPUs / TPUs.
As researchers begin to collect larger in-situ multiplexed imaging datasets, there is a growing need for more scalable approaches for analysis of microenvironments.
Here we propose a
VAE-style network which we call Neural Spatial LDA extending the auto-encoding
Variational Bayes formulation of classical LDA from [18].
We show Neural Spatial LDA achives significant speed-up over Spatial LDA while at the same time recovering similar topic distributions thus enabling its use in large data domains.
Introduction
The development of in-situ multiplexed profiling techniques, has enabled high-dimensional measurements of the abundance of protein & expression of genes while preserving spatial information [12, 4, 19, 17, 16, 11, 6].
For example, Keren et al.
[8] characterized tumor-immune microenvironments in TNBC using Multiplixed Ion Bean Imaging by Time Of Flight (MIBI-ToF) [9] and categorized microenvironments into three types based on the spatial organization of tumor and immune cells and further demonstrated that these microenvironments stratified patient survival in their cohort.
Similarly,
Mccaffrey et al.
[13] identified specific immuno-suppressive microenvironments in Tuberculosis granulomas which were associated with patient outcomes.
Both examples demonstrate the potential for spatial analysis to yield prognostically relevant markers of disease.
While there exist various techniques to delineate cellular microenvironments [8, 15], in our work we focus on accelerating Spatial Latent Dirichlet Allocation (Spatial LDA) [5] which models cellular microenvironments in a probabilistic manner using topic models that incorporate a spatial prior - that neighboring microenvironments are more likely to have similar distribution of cell types.
The current implementation of Spatial LDA uses mean field variational inference for determining the posterior distribution of the parameters of the generative model and does not leverage modern accelerators like
GPUs or TPUs.
As a result it takes hours to fit a model to even moderately sized datasets today.
In this work, we propose an encoder-decoder style network based on auto-encoding Variational Bayes
LDA 18 with spatial regularization of discovered topics, which we call Neural Spatial LDA .
Our implementation of Neural Spatial LDA runs on GPUs / TPUs and we show on several previously published datasets that the microenvironments detected by Neural Spatial LDA are similar to Spatial
LDA but with significantly faster training and inference times.
Finally, reformulating Spatial LDA as Neural Spatial LDA also opens the possibility of further extensions to the Spatial LDA family of models, for example enabling supervised variants to specifically find microenvironments that are predictive of patient outcomes or regional annotations, extensions that we are currently exploring.
Neural Spatial LDA
As in Spatial LDA[5], we cast the identification of micro-environments from multiplexed images as a topic modeling problem, where a microenvironment is considered a document , unordered set of cells that belong to a microenvironment are analogous to words in a document .
The task then is to identify both distribution cells / words per topic and distribution of topics for each microenvironment / document using cell counts.
First, we briefly review LDA [2] and Spatial LDA, we then describe the auto-encoding Variational Bayes version of Spatial LDA.
Spatial LDA assumes the data to have a generative process as follows,
[Figure]
Let us consider each in-situ profiled image to be a graph G with vertices V = { v i , ...v N } representing microenvironments and edges E with each e ij ∈ E denoting an edge from vertex v i to v j .
Spatial
LDA adds a Laplacian constraint over every pair of α i & α j of microenvironments v i & v j that are connected by an edge e ij and introduces a hyperparameter σ controling the strength of this prior over topic distributions of microenvironments.
Intuitively, the graph G defines neighborhoods that should have similar topic distributions and σ controls how strongly that inductive bias is enforced.
The rest of the generative process is a standard LDA model with each cell w n being sampled by first choosing a topic distribution θ k from a Dirichlet prior parameterized by α , then choosing a topic z n and then further choosing a cell from the categorical distribution of cells parameterized by β z n .
Here, K represents the number of topics and V represents the size of the vocabulary.
Note that, in the limit where the graph G is a fully connected graph (i.e.) E = { e ij | ∀
( i, j ) ∈ { 1 , ..., N } , i ̸
= j } and σ → 0 , all α ’s are constrained to be the same and Spatial LDA reduces to classical LDA.
In that scenario, all microenvironments are assumed to have the same cell distributions and spatial information is lost.
In order to capture spatial structure in multiplexed images Spatial
LDA defines a graph that captures spatial relationships between local neighborhood of cells.
In the original paper, the authors suggest defining E by taking the Voronoi partitioning of neighborhood positions and connecting neighborhoods that share a facet in Voronoi partitions.
Spatial LDA then uses mean field variational inference for learning the parameters of the model by optimizing the Evidence Lower BOund (ELBO) shown in (1) where γ, ϕ, ξ are variational parameters for approximating distributions parametrized by θ , z , α respectively. L
[Equation]
Owing to the conjugacy of categorical and Dirichlet distributions, the coordinate update steps have a closed form in classical LDA, while Spatial LDA involves optimizing a non-smooth function for the variational parameter ξ , due to the presence of the Laplacian prior.
This was overcome by using alternating direction method of multipliers (ADMM)
+ primal-dual interior point optimization approaches (we refer the reader to Spatial LDA[5] for the exact update steps).
We observe that these optimization steps do not scale well with the number of data points compared to gradient based methods for training modern deep neural networks.
In addition, parameter updates of Spatial LDA are not optimized to run on GPUs / TPUs.
Therefore, to enable faster training and inference, we propose to use auto-encoding Variational Bayes (AEVB)[10] for Spatial LDA.
Srivastava et al.
[18] proposed two key ingredients that made it possible to apply the re-parametrization trick to classical LDA.
Specifically, 1) Marginalize out the discrete variables z ’s 2) Replace Dirichlet prior with Logistic Normal with diagonal covariance in softmax basis.
Logistic normal could be easily re-parametrized by sampling from the standard normal as, E z ∼LN ( µ , Σ ) [ f ( z )] =
E ϵ
∼N (0 ,I )
[ f ( ρ ( µ + Σ 12 ϵ ))] where ρ represents the softmax function.
Neural Spatial LDA uses these two ideas along with spatial regularization as given below,
2.1 Spatial Regularization
The final piece in the puzzle and our main contribution is to enforce spatial regularization on topic distribution of microenvironments.
This could be done through two different ways, 1) sample α ’s such that they obey p ( α ) ∝ (cid:81)
e ij ∈E Laplace
( α i − α j , σ ) and compute prior p ( θ | α , σ ) =
LN ( θ | α , σ ; µ , Σ ) that will impose the constraint through the KL-term of the ELBO or 2) apply the Laplacian constraint directly on µ ’s. Since VAEs learn variational posteriors are far away from the original priors they start with after training, we opted to apply L1 constraint on µ ’s of microenvironments that are connected via an edge in G .
The final objective function for our proposed
Neural Spatial LDA is the following,
[Equation]
Experiments
We applied Neural Spatial LDA to three different datasets - 1)
Simulated: A simulated dataset with known ground truth topic distribution and "cell" distribution under each topic (details in Appendix).
2) Mouse spleen:
A dataset of CODEX images of mouse spleen with a panel of 30 different antibodies from [7]. 3) TB:
A dataset of MIBI images of granulomas in Tuberculosis patients from [13].
As in the original Spatial LDA paper, we preprocess CODEX and MIBI-ToF multichannel images to segment cells using DeepCell [1].
Then we get mean expression profiles of each cell and then cluster these profiles using FlowSOM [20].
We then associate cell index to each cluster and then get the count of different types of cells in a spatial neighborhood of each cell.
The cell counts of different cell types around a spatial neighborhood constitute a microenvironment and we use them as input features.
To qualitatively evaluate whether Spatial LDA and Neural Spatial LDA infer similar topic distributions to various microenvironments, we run Spatial LDA, & Neural Spatial LDA on all three datasets and color each cell by its most probable topic.
From Fig.[3], we observe that Spatial LDA and Neural
Spatial LDA groups similar regions of cells into similar topics across all three datasets.
(note: the topics are matched manually after training since the order of the topics is not guaranteed to match)
To quantitatively measure the similarity between topic distributions learnt by both methods, we randomly sample pairs of cells from each image without replacement, compute the KL-divergence between the topic distributions of pairs of cells and report the Spearman correlation between Spatial
LDA and Neural Spatial LDA.
Our intuition is that this KL divergence should capture the similarity between 2 neighborhoods under each model and a high correlation indicates that Spatial and Neural
Spatial LDA both capture similar notions of similarity / differences in neighborhoods.
This also allows us to evaluate the agreement between Spatial and Neural Spatial LDA in a way that is agnostic to reordering of topics.
We see from Fig.[2], that Spatial LDA and Neural Spatial LDA yields highly correlated KL-divergences between topic distributions for the same pairs of cells indicating good agreement between the two methods.
We note that we expect some differences in KL-divergence due to the approximations described above - marginalizing z ′ s & replacing Dirichlet with Logistic
Normal, as well as differences - in optimization co-ordinate ascent vs gradient descent - used for inference.
Next, we measure the time it takes for Spatial LDA and Neural Spatial LDA to train using the same machine, with Spatial LDA was run on the CPU as it was not optimized to run on GPUs while
Neural Spatial LDA was run a single Titan RTX.
The stopping criteria for Spatial LDA is set to be the same number of iterations used in the original work while Neural Spatial LDA was trained til convergence of the ELBO.
From Fig.[4], we could see that AEVB version of Spatial LDA achieves ∼ 84%, ∼ 95% & ∼ 90% reduction in training times on Mouse spleen, Simulated & Tuberculosis datasets respectively.
Conclusion
Spatial LDA is a probabilistic model used to identify TMEs in an unsupervised way but its inability to utilize GPUs / TPUs and slow training and inference limits its use in large data domain applications.
We propose Neural Spatial LDA and show on previously published datasets that we recover similar microenvironments while acclerating training and inference significantly.
A Simulated Dataset
We create a simulated dataset according to the p ( θ ) and p ( β ) shown in (2) & (3), given ( x, y ) co-ordinate of a "cell".
[Equation]
We set H = 32 , W = 32 & V = 4 for our experiments on this simulated dataset.
Fig.[5] shows the topic distribution for a 32 × 32 grid of cells.