Graph Topic Modeling for Documents with Spatial or Covariate Dependencies

In this paper, we propose the first pLSI method that can be made amenable to the inclusion of additional information on the similarity between documents, as encoded by a known graph. More specifically:

• a.

  We propose a scalable algorithm based on a graph-aligned singular decomposition of the empirical frequency matrix $X$ to provide estimates of $W$ and $A$ (Section 2). Additionally, we develop a new cross-validation procedure for our graph-based penalty that allows us to choose the optimal regularization parameter adaptively (Section A of the Appendix).

• b.

  We prove the benefits of the graph alignment procedure by deriving high probability upper bounds for both $W$ and $A$ in Section 3, which we verify through extensive simulations in Section 3.4.

• c.

  Finally, we showcase the advantage of our method over LDA-based methods and non-structured pLSI techniques on three real-world datasets: two spatial transcriptomics examples and a recipe dataset in Section 4.

Throughout this paper, we use the following notations. For any $t\in\mathbb{Z}_{+}$, $[t]$ denotes the set $\{1,2,...,t\}$. For any $a,b\in\mathbb{R}$, we write $a\vee b=\max(a,b)$ and ${a\wedge b=\min(a,b)}$. We use $\mathbf{1}_{d}\in\mathbb{R}^{d}$ to denote the vector with all entries equal to 1 and $\mathbf{e}_{k}\in\mathbb{R}^{d}$ to denote the vector with $k^{th}$ element set to 1 and 0 otherwise. For any vector $u$, its $\ell_{2}$, $\ell_{1}$ and $\ell_{0}$ norms are defined respectively as $\|u\|_{2}=\sqrt{\sum_{i}u_{i}^{2}}$, $\|u\|_{1}=\sum_{i}|u_{i}|$, and $\|u\|_{0}=\sum_{i}\mathbf{1}\{u_{i}\neq 0\}$. Let $I_{m}$ denote the $m\times m$ identity matrix. For any matrix ${A=(a_{ij})\in\mathbb{R}^{n\times p}}$, $A(i,j)$ denote the $(i,j)$-entry of $A$, $A_{i\cdot}$ and $A_{\cdot j}$ denote the $i^{\text{th}}$ row and $j^{th}$ column of $A$ respectively. Throughout this paper, $\lambda_{i}(A)$ stands for the $i^{th}$ largest singular value of the matrix $A$ with ${\lambda_{\max}(A)=\lambda_{1}(A)}$, ${\lambda_{\min}(A)=\lambda_{p\wedge\text{rank}(A)}(A)}$. We also denote as $U_{K}(A)$ and $V_{K}(A)$ the left and right singular matrix from the rank-$K$ SVD of $A$. The Frobenius, entry-wise $\ell_{1}$ norm and the operator norms of $A$ are denoted as $\left\|A\right\|_{F}=\sqrt{\sum_{i,j}a_{ij}^{2}}$, $\lVert A\rVert_{11}=\sum_{i,j}|a_{ij}|$, and ${\left\|A\right\|_{op}=\lambda_{1}(A)}$, respectively. The $\ell_{21}$ norm is denoted as $\big{\|}A\big{\|}_{21}=\sum_{i}\|A_{i\cdot}\|_{2}$. For any positive semi-definite matrix $A$, $A^{1/2}$ denotes its principal square root that is positive semi-definite and satisfies $A^{1/2}A^{1/2}=A$. The trace inner product of two matrices $A,B\in\mathbb{R}^{n\times p}$ is denoted by $\langle A,B\rangle=\mathop{\sf Tr}(A^{\top}B)$. $A^{\dagger}$ denotes the pseudo-inverse of the matrix $A$ and $P_{A}$ denotes the projection matrix onto the subspace spanned by columns of $A$.

Let $\mathcal{G}=(\mathcal{D},\mathcal{E})$ denote an undirected graph induced by a known adjacency matrix on the $n$ documents in the corpus. The documents are represented as nodes $\mathcal{D}$, and $\mathcal{E}$ denotes the edge set of size $\lvert\mathcal{E}\rvert=m$. Throughout this paper, for simplicity, we will assume that $\mathcal{G}$ is binary, but our approach—as discussed in Section 2.3—can be in principle extended to weighted graphs. We denote the graph’s incidence matrix as $\Gamma\in\mathbb{R}^{m\times n}$ where, for any edge $e=(i,j),i<j$ between nodes $i$ and $j$ in the graph, $\Gamma_{ei}=1$, $\Gamma_{ej}=-1$ and $\Gamma_{ek}=0$ for any $k\neq i,j.$ It is easy to show that the Laplacian of the graph can be expressed in terms of the incidence matrix as $L=\Gamma^{\top}\Gamma$ (Hütter & Rigollet 2016). Let $\Gamma^{\dagger}$ be the pseudo-inverse of $\Gamma$, and denote by $\mathbf{s}_{i},i=1\cdots m$ its columns, so that $\Gamma^{\dagger}=[\mathbf{s}_{1},\cdots,\mathbf{s}_{m}]$. We also define the inverse scaling factor of the incidence matrix $\Gamma$ (Hütter & Rigollet 2016), a quantity necessary for assessing the performance of GpLSI:

We focus on the estimation of the topic mixture matrix under the assumption that neighboring documents have similar topic mixture proportions: $W_{i\cdot}\approx W_{j\cdot}$ if $i\sim j$. This implies that the rows of $W$ are assumed to be smooth with respect to the graph $\mathcal{G}$. Noting that the difference of mixture proportions between neighboring nodes $i$ and $j$ ($e=(i,j)\in\mathcal{E}$) can be written as $(\Gamma W)_{d\cdot}=W_{i\cdot}-W_{j\cdot}$, this smoothness assumption effectively implies sparsity on the rows of the matrix $\Gamma W$.

The previous assumption is akin to assuming that the underlying mixture matrix $W$ is piecewise-continuous with respect to the graph $\mathcal{G}$, or more generally, that it can be well approximated by a piecewise-continuous function.

Our setting is not limited to connected document graphs. Denote $n_{\mathcal{C}}$ the number of connected subgraphs of $\mathcal{G}$ and $n_{\mathcal{C}_{l}}$ the number of nodes in the $l^{th}$ connected subgraph. Let $n_{\mathcal{C}_{\min}}$ be the size of the smallest connected component:

The error bound of our estimators will depend on both $n_{\mathcal{C}}$ and $n_{\mathcal{C}_{\min}}$. In the rest of this paper, we will assume that all connected components have roughly the same size: $n_{\mathcal{C}_{1}}\asymp\dots\asymp n_{\mathcal{C}_{n_{\mathcal{C}}}}$.

Since the smoothness assumption (Assumption 1) pertains to the rows of the document-topic mixture matrix $W$, we build on the pLSI algorithm developed by Klopp et al. (2021). In this subsection, we provide a brief overview of their method.

When we assume we directly observe the true expected frequency matrix $M$ defined in Equation (2), Klopp et al. (2021) propose a fast and simple procedure to recover the mixture matrix $W$. Specifically, let $U\in\mathbb{R}^{n\times K}$ and $V\in\mathbb{R}^{p\times K}$ be the left and right singular vectors obtained from the singular value decomposition (SVD) of the true signal $M\in\mathbb{R}^{n\times p}$, so that $M=U\Lambda V^{T}$. A critical insight from their work is that $U$ can be decomposed as:

where $H$ is a full-rank, $K\times K$-dimensional matrix. From this decomposition, it follows that the rows of $U$, which can be viewed as $K$-dimensional embeddings of the documents, lie on the $K$-dimensional simplex $\Delta_{K-1}$. The simplex’s vertices, represented by the rows of $H$, correspond to the anchor documents (Assumption 2). Thus, identifying these vertices through any standard vertex-finding algorithm applied to the rows of $U$ will enable the estimation of $W$. The procedure of Klopp et al. (2021) can be summarized as follows:

Step 1:

Compute the singular value decomposition (SVD) of the matrix $M$, reduced to rank $K$, to obtain: $M=U\Lambda V^{T}$.

Step 2: Vertex-Hunting Step:

Apply the successive projection algorithm (SPA) (Araújo et al. 2001) (a vertex-hunting algorithm) on the rows of $U$. This algorithm returns the indices of the selected “anchor documents,” $J\subseteq[n]$ with $|J|=K$. Define $\widehat{H}=U_{J\cdot}$, where each row corresponds to one of the $K$ vertices of the simplex $\Delta_{K-1}$.

Step 3: Recovery of $W$:

$W$ can simply be recovered from $U$ and $\widehat{H}$ as

$$\widehat{W}={U}\widehat{H}^{-1}.$$

(7)

Step 4: Recovery of $A$:

Finally, $A$ can subsequently be estimated as $\widehat{A}=\widehat{H}\Lambda V^{\top}.$

In the noisy setting, the procedure is adapted by plugging the observed frequency $X$ instead of $M$ in Step 1 and getting estimates of the singular vectors: $X=\widehat{U}\widehat{\Lambda}\widehat{V}^{\top}$. Under a similar set of assumptions as ours (Assumptions 2-4), Theorem 1 of (Klopp et al. 2021) states that the error of $\widehat{W}$ is such that: $\min_{P\in\mathcal{P}}\|\widehat{W}-WP\|_{F}\leq{\rm C_{0}}K\sqrt{n\log(n+p)/N}$, where $\mathcal{P}$ denotes the set of all permutation matrices. Their analysis provides one of the best error bounds on the estimation of the topic mixture matrix $W$ for pLSI.

However, this approach has two key limitations. First, the consistency of their estimator relies on having a sufficiently large number of words per document, $N$. In particular, a necessary condition for the aforementioned results to hold is that $N\geq CK^{5}\log(n+p)$. The authors establish minimax error bounds, showing that the rate of any estimator’s error for $W$ is bounded below by a term on the order of $O(\sqrt{n/N})$ (Theorem 3, Klopp et al. (2021)). In other words, the accurate estimation of $W$ requires that each document contains enough words. In many practical scenarios—such as the tumor microenvironment example mentioned earlier—this condition may not hold. However, we might still have access to additional information indicating that certain documents are more similar to one another.

Second, the method is relatively rigid and does not easily accommodate additional structural information, such as document-level similarities. Indeed, this method does not rely on a convex optimization formulation to which we could simply add a regularization term, and the vertex-hunting algorithm does not readily incorporate metadata of the documents.

Theoretical insights from Klopp et al. (2021) help explain why topic modeling deteriorates in low-$N$ regimes. When the number of words per document is too small, the observed frequency matrix $X$ can be viewed as a highly noisy approximation of $M$, causing the estimated singular vectors $\widehat{U}$ to deviate significantly from the true underlying point cloud $U$. To mitigate this issue, Klopp et al. (2021) suggest a preconditioning step that improves the estimation of the singular vectors in noisy settings.

In this paper, we take a different approach by exploiting the graph structure associated with the documents to reduce the noise in $X$. Rather than preconditioning the empirical frequency matrix, we propose an additional denoising step that leverages the graph structure to produce more accurate estimates of $U$, $V$, and $\Lambda$. Specifically, we modify the SVD of $X$ in Step 1 and estimation of topic matrix $A$ in Step 4 described in Section 2.2.

Step 1: Iterative Graph-Aligned SVD of $X$:

We replace Step 1 of Section 2.2 with a graph-aligned SVD of the empirical word-frequency matrix $X$. More specifically, in the graph-aligned setting, we assume that the underlying frequency matrix $M$ belongs to the set:

$$\begin{split}\mathcal{F}(n,p,K,s)=\{&M=U\Lambda V^{\top}\in\mathbb{R}^{n\times p},\ \mathrm{rank}(M)=K:\\ &U\in\mathbb{R}^{n\times K},V\in\mathbb{R}^{p\times K},\Lambda=\mathrm{diag}(\lambda_{1},\lambda_{2},\cdots,\lambda_{K}),\\ &|\mathrm{supp}(\Gamma U)|\leq s,\lambda_{K}>0\}.\end{split}$$

(8)

Throughout the paper, we shall allow $s,p,N,$ and $n$ to vary. We will assume the number of topics $K$ to be fixed.

Step 2, 3

Same as Step 2,3 in Section 2.2.

Step 4: Recovery of $A$:

$A$ can subsequently be estimated by solving a constrained regression problem of $X$ on $\widehat{W}$:

$$\begin{split}\widehat{A}&=\text{argmin}_{A\in\mathbb{R}^{K\times p}}\|X-\widehat{W}A\|_{F}^{2}\\ &\text{ such that }\forall k\in[K],\ \sum_{j=1}^{p}A_{kj}=1,\ A_{kj}\geq 0\end{split}$$

(9)

The improvement in the estimation of the singular vectors induced by our iterative denoising procedure is characterized in the following theorem.

The proof of this result is provided in Section C of the Appendix.

We now show how our denoised singular vectors can be used to improve the estimation of the mixture matrix $W.$

Theorem 3 shows that $\widehat{W}$ is highly accurate as long as the document lengths are large enough, as defined by $N\gtrsim({n_{\mathcal{C}}}+\rho^{2}(\Gamma){s\lambda_{\max}(\Gamma)})\log(n)/n$. This requirement is more relaxed than the condition $N\gtrsim\log(n+p)$ for pLSI provisioned in Theorem 1 and Corollary 1 of Klopp et al. (2021). This indicates that GpLSI can produce accurate estimates even for smaller $N$, by sharing information among neighboring documents. The shrinkage of error due to graph-alignment is characterized by the term $\frac{1}{\sqrt{n}}(\sqrt{n_{\mathcal{C}}}+\rho(\Gamma)\sqrt{s\lambda_{\max}(\Gamma)})$, which is equal to one when the graph $\mathcal{G}$ is empty. In general, the effect of the regularization depends on the graph topology. Hütter & Rigollet (2016) show in fact that the inverse scaling factor verifies: $\rho(\Gamma)\leq\sqrt{2}/\sqrt{\lambda_{n-1}(L)}$. The quantity $\lambda_{n-1}(L)$, also known as the algebraic connectivity, provides insights into the properties of the graph, such as its connectivity. Intuitively, higher values of $\lambda_{n-1}(L)$ reflect more tightly connected graphs (Chung 1997), thereby reducing the effective degrees of freedom induced by the graph-total variation penalty. By contrast, $\lambda_{\max}(\Gamma)$ can be coarsely bounded using the maximum degree $d_{\max}$ of the graph: $\lambda_{\max}(\Gamma)\leq\sqrt{2d_{\max}}$ (Anderson Jr & Morley 1985, Zhang 2011). Consequently, we can expect our procedure to work well on well-connected graphs with bounded degree. Examples include for instance grid-graphs, $k$-nearest neighbor graphs, or spatial (or planar) graphs. We provide a more detailed discussion and more explicit bounds for specific graph topologies in Section 3.3.

Furthermore, using the inequality $\|\widehat{W}-WP\|_{11}\leq\sqrt{Kn}\|\widehat{W}-WP\|_{F}$, it immediately follows that:

Finally, we characterize the error bound of $\widehat{A}$. The full proofs of Theorems 3 and 4 are deferred to Section D of the Appendix.

We now analyze the behavior of the error bound provided in Theorem 3 for different graph structures, further expliciting the dependency of our bounds on graph properties.

We evaluate the performance of our method using synthetic datasets where $W$ is aligned with respect to a known graph.

We first consider the analysis of CODEX data, which allows the identification of individual cells in tissue samples, holding valuable insights into cell interaction profiles, particularly in cancer, where these interactions are crucial for immune response. Since cellular interactions are hypothesized to be local, these patterns are often referred to as “tumor microenvironments”. In the context of topic modeling, we can regard a tumor microenvironment as a document, immune cell types as words, and latent characteristics of a microenvironment as a topic. However, due to the small number of words per document, the recovery of the topic mixture matrix and the topics themselves can prove challenging. Chen et al. (2020) propose using the adjacency of documents to assign similar topic proportions to neighboring tumor cells. Similarly, we construct a spatial graph based on proximity of tumor microenvironments to uncover novel tumor-immune cell interaction patterns.

The first CODEX dataset is a collection of 292 tissue samples from 161 colorectal cancer patients collected at Stanford University (Wu et al. 2022). The locations of the cells were retrieved using a Voronoi partitioning of the sample, and the corresponding spatial graphs were constructed encoding the distance between microenvironments. More specifically, we define a tumor microenvironment as the 3-hop neighborhood of each cell, following the definition originally used by Wu et al. (2022). Each microenvironment contains 10 to 30 immune cells of 8 possible types. This aligns with the setting where the document length $N<30$ is small compared to the vocabulary size $p=8$.

We aggregate frequency matrices and tumor cellular graphs of all samples and fit three methods: our proposed GpLSI, the original pLSI approach of Klopp et al. (2021), and LDA (Blei et al. 2001) to estimate tumor-immune topics. The estimated topic weights of $K=6$ are illustrated in Figure 3(A). After aligning topic weights across methods, we observe similar immune topics (Topic 1, 2, 3) in GpLSI and LDA which are not found in the estimated topics of pLSI.

To determine the optimal number of topics $K$, we use the method proposed by Fukuyama et al. (2023). In this work, the authors construct “topic paths” to track how individual topics evolve, split or merge, as the number of topics $K$, increases. We observe in Figure 3(B) that while the GpLSI path has non-overlapping topics up to $K=6$, other methods fail to provide consistent and well-separated topics.

To evaluate the quality and stability of the recovered topics, we also measure the coherence of the estimated topic weights of batches of samples, as suggested in Tran et al. (2023). We divide 292 samples into five batches and estimate the topic weights $A^{b}$ for $b\in[5]$. For every pair of batches $(b,b^{{}^{\prime}})$, we align $A^{b}$ and $A^{b^{{}^{\prime}}}$ (we permute $A^{b^{{}^{\prime}}}$ with $P$ where $P=\arg\min_{P\in\mathcal{P}}\|A^{b}-PA^{b^{{}^{\prime}}}\|$) and measure the entry-wise $\ell_{1}$ distance and cosine similarity. We repeat this procedure five times and plot the scores in Figure 3(C). We notice that GpLSI provides the most coherent topics across batches for $K=3,4,5$. Combining with the metrics of LDA, we can choose the optimal $K$ as 5 or 6.

Next, we conduct survival analysis to identify the immune topics associated with higher risk of cancer recurrence. We consider two logistic models with different covariates to predict cancer recurrence and calculate the area under the curve (AUC) of the receiver-operating characteristic (ROC) curves to evaluate model performance.

In the first model, we use the proportion of each microenvironment topic as covariates for each sample. Since the $K$ covariates sum up to one, we apply isometric log-ratio transformation to represent it with $K-1$ orthonormal basis vectors In the second model, we dichotomize each topic proportion to low and high proportion groups. The cutoffs are determined using the maximally selected rank statistics.

The AUC for each number of topics is shown in Figure 4(A). GpLSI achieves the highest area under the curve (AUC) at $K=6$ in both models. We also plot Kaplan Meier curves for each topic using the same dichotomized topic proportions. The result for GpLSI is illustrated in Figure 4(B). We observe that Topic 2, which is characterized by a high prevalence of granulocytes, and Topic 6, a mixture of CD4 T cells and blood vessels, are associated with lower cancer recurrence. Positive effect of granulocytes on cancer prognosis was also reported by Wu et al. (2022), who found out that a microenvironment with clustered granulocyte and tumor cells is associated with better patient outcomes. We observe the same association of granulocyte with lower risk in LDA Figure 16 of the Appendix.

We also apply our method to identify immune topics in mouse spleen. In this setting, each document is anchored to a B cell (Chen et al. 2020). A previous study has processed the original CODEX images from Goltsev et al. (2018) to obtain the frequencies of non-B cells in the 100 pixel neighborhood of each B cell (Chen et al. 2020). The final input for the topic models consists of a 35,271 B cell microenvironments by 24 cell types frequency matrix, along with the positional data of B cells.

In this example, we evaluate GpLSI by examining whether the introduction of our graph-based regularization term in the estimation of topic mixture matrices enhances document clustering. Figure 5(A) presents the estimated topics for all models at $K=3,5,7,10$. Notably, the topics derived from GpLSI, pLSI, and LDA more clearly demarcate distinct B cell microenvironment domains compared to those estimated by TopicSCORE and LDA. Among these three methods, GpLSI yields the least noisy cellular clustering, as evidenced by the magnified view of a selected subdomain.

We also evaluate the quality of clusters with two metrics, Moran’s I and the percentage of abnormal spots (PAS) (Shang & Zhou 2022). Moran’s I is a classical measure of spatial autocorrelation that assesses the degree to how values are clustered or dispersed across a spatial domain. PAS score measures the percentage of B cells for which more than 60% of its neighboring B cells have different topics. Higher Moran’s I and lower PAS score indicate more spatial smoothness of the estimated topics. From Figure 5(B), we conclude that GpLSI has the highest Moran I, and the lowest PAS scores, demonstrating improved spatial smoothness of the topics.

We observe that the B cell microenvironment topics identified with GpLSI align well with their biological context (Figure 5(C)). By referencing the manual annotations of B cells from the original study by Goltsev et al. (2018), we infer that Topic 1, Topic 2, Topic 3, and Topic 5 correspond to the red pulp, periarteriolar lymphoid sheath (PALS), B-follicles, and the marginal zone. This interpretation is further supported by high expression of CD4+T cells in Topic 2 (PALS) and high expression of marginal zone macrophages in Topic 5 (marginal zone).

This dataset contains recipes from 20 different cuisines across Europe, Asia, and South America. Each recipe is a list of ingredients which allow us to convert to a count matrix with 13,597 recipes (documents) and 1,019 unique ingredients (words). Under the assumption that neighboring countries would have similar cuisine styles, we construct a graph of recipes based on the geographical proximity of the countries. Specifically, for each recipe, we select the five closest recipes from neighboring countries (including its own country) based on the $\ell_{1}$ distance of the ingredient count vectors and define them as neighboring nodes on the graph. Through this, we aim to identify general cooking styles that are prevalent across various countries worldwide.

We run GpLSI, pLSI, and LDA with $K=5,7,10,15,20$ topics. We illustrate the estimated topics of GpLSI for $K=7$ in Figure 6. The results for pLSI and LDA are provided in Section G of the Appendix. With this approach, Topic 1 is clearly a baking topic and Topic 6 is defined by strong spices and sauces common in Mexican or parts of Southeast Asian cuisines. We also observe a general topic for Asian cuisines (Topic 2) and another for Western countries (Topic 7). To evaluate the estimated topics, we compare each topic’s characteristics with the cuisine-by-topic proportion (Figure 6(C)). Indeed, the style of each topic defined by the anchor ingredients aligns with the cuisines that have a high proportion of that topic. For example, the baking topic (Topic 1) is prevalent in British, Irish, French, Russian, and South American cuisines.

In contrast, for pLSI, it is difficult to analyze the characteristics for each topic because Topics 1-4 have one or no identified anchor ingredients. Comparing the cuisine-by-topic proportions of GpLSI and LDA, we observe that GpLSI reveals many cuisines as mixtures of different cooking styles (Figure 6(B)). In contrast, for LDA, many cuisines such as Moroccan, Mexican, Korean, Chinese, Thai have their recipes predominantly classified to a single topic (Figure 18(B) of the Appendix). GpLSI provides estimates of topic mixture and topic weights that are more relevant to our goal of discovering global cooking styles.