Estimation and inference for the Wasserstein distance between mixing measures in topic models

Given a set $\mathcal{A}\subseteq\mathcal{P}(\mathcal{Y})$ of probability distributions on a Polish space $\mathcal{Y}$, consider the set $\mathcal{S}$ of finite mixtures of elements of $\mathcal{A}$. As long as mixtures consisting of elements of $\mathcal{A}$ are identifiable, we may associate any $r^{(1)}=\sum_{k=1}^{K}\alpha_{k}^{(1)}A_{k}\in\mathcal{S}$ with a unique mixing measure $\boldsymbol{\alpha}^{(1)}=\sum_{k=1}^{K}\alpha_{k}^{(1)}\delta_{A_{k}}$, which is a probability measure on $\mathcal{A}\subseteq\mathcal{Y}$.

We assume that $\mathcal{P}(\mathcal{Y})$ is equipped with a metric $d$. This may be one of the classic distances on probability measures, such as the Total Variation (TV) or Hellinger distance, or it may be a more exotic distance, for instance, one computed via a machine learning procedure or tailored to a particular data analysis task. As mentioned above, given mixing measures $\boldsymbol{\alpha}^{(1)}$ and $\boldsymbol{\alpha}^{(2)}$, corresponding to mixture distributions $r^{(1)}$ and $r^{(2)}$, respectively, it has become common in the literature to compare them by computing their Wasserstein distance $W(\boldsymbol{\alpha}^{(1)},\boldsymbol{\alpha}^{(2)};d)$ with respect to the distance $d$ (see, e.g., [49, 31]). It is intuitively clear that computing a distance between the mixing measures will indicate, in broad terms, the similarity or dissimilarity between the mixture distributions; however, it is not clear whether using the Wasserstein distance between the mixing measures best serves this goal, nor whether such a comparison has any connection to a bona fide distance between the mixture distributions themselves.

As our first main contribution, we show that, indeed, the Wassersetin distance between mixing measures is a fundamental quantity to consider and study, but we arrive at it from a different perspective. In Section 3 we introduce a new distance, the Sketched Wasserstein Distance ($\operatorname{SWD}$) on $\mathcal{S}$, between mixture distributions which is defined and uniquely characterized as the largest—that is, most discriminative—jointly convex metric on $\mathcal{S}$ for which the inclusion $\iota:(\mathcal{A},d)\to(\mathcal{S},\operatorname{SWD})$ is an isometry. Despite the abstractness of this definition, we show, in Theorem 1 below that

This fact justifies our name Sketched Wasserstein Distance for the new metric—it is a distance on mixture distributions that reduces to a Wasserstein distance between mixing measures, which may be viewed as “sketches” of the mixtures. We use the results of Section 3 as further motivation for developing inferential tools for the Wasserstein distance between mixing measures, in the specific setting provided by topic models. We present this below.

To state our results, we begin by introducing the topic model, as well as notation that will be used throughout the paper. We also give in this section relevant existing estimation results for the model. They serve as partial motivation for the results of this paper, summarized in the following sub-sections.

Topic models are widely used modeling tools for linking $n$ discrete, typically high dimensional, probability vectors. Their motivation and name stems from text analysis, where they were first introduced [13], but their applications go beyond that, for instance to biology [15, 18]. Adopting a familiar jargon, topic models are used to model a corpus of $n$ documents, each assumed to have been generated from a maximum of $K$ topics. Each document $i\in[n]:=\{1,\ldots,n\}$ is modeled as a set of $N_{i}$ words drawn from a discrete distribution $r^{(i)}\in\Delta_{p}:=\{v\in\mathbb{R}_{+}^{p}:v^{\top}\mathds{1}_{p}=1\}$ supported on the set of words $\mathcal{Y}:=\{y_{1},\ldots,y_{p}\}$, where $p$ is the dictionary size. One observes $n$ independent, $p$-dimensional word counts $Y^{(i)}$ for $i\in[n]$ and assumes that

For future reference, we let $X^{(i)}:={Y^{(i)}/N_{i}}$ denote the associated word frequency vector.

The topic model assumption is that the $p\times n$ matrix $R$, with columns $r^{(i)}$, the expected word frequencies of each document, can be factorized as $R=AT$, where the $p\times K$ matrix $A:=[A_{1},\ldots,A_{K}]$ has columns $A_{k}$ in the probability simplex $\Delta_{p}$, corresponding to the conditional probabilities of word occurrences, given a topic. The $K\times n$ matrix $T:=[\alpha^{(1)},\ldots,\alpha^{(n)}]$ collects probability vectors $\alpha^{(i)}\in\Delta_{K}$ with entries corresponding to the probability with which each topic $k\in[K]$ is covered by document $i$, for each $i\in[n]$. The only assumption made in this factorization is that $A$ is common to the corpus, and therefore document independent: otherwise the matrix factorization above always holds, by a standard application of Bayes’s theorem. Writing out the matrix factorization one column $i\in[n]$ at a time

one recognizes that each (word) probability vector $r^{(i)}\in\Delta_{p}$ is a discrete mixture of the (conditional) probability vectors $A_{k}\in\Delta_{p}$, with potentially sparse mixture weights $\alpha^{(i)}$, as not all topics are expected to be covered by a single document $i$. Each probability vector $r^{(i)}$ can be written uniquely as in (2) if uniqueness of the topic model factorization holds. The latter has been the subject of extensive study and is closely connected to the uniqueness of non-negative matrix factorization. Perhaps the most popular identifiability assumption for topic models, owing in part to its interpretability and constructive nature, is the so called anchor word assumption [25, 4], coupled with the assumption that the matrix $T\in\mathbb{R}^{K\times n}$ has rank $K$. We state these assumptions formally in Section 4.1. For the remainder of the paper, we assume that we can write $r^{(i)}$ uniquely as in (2), and refer to $\alpha^{(i)}$ as the true mixture weights, and to $A_{k}$, $k\in[K]$, as the mixture components. Then the collection of mixture components defined in the previous section becomes $\mathcal{A}=\{A_{1},\dots,A_{K}\}$ and the mixing measure giving rise to $r^{(i)}$ is

which we view as a probability measure on the metric space $\mathcal{X}:=(\Delta_{p},d)$ obtained by equipping $\Delta_{p}$ with a metric $d$.

Computationally tractable estimators $\widehat{A}$ of the mixture components collected in $A$ have been proposed by [4, 3, 61, 54, 9, 10, 67, 38], with finite sample performance, including minimax adaptive rates relative to $\|\widehat{A}-A\|_{1,\infty}=\max_{1\leq k\leq K}\|\widehat{A}_{k}-A_{k}\|_{1}$ and $\|\widehat{A}-A\|_{1}=\sum_{k=1}^{K}\sum_{j=1}^{p}|\widehat{A}_{jk}-A_{jk}|$ obtained by [61, 9, 10]. In Section 4.1 below we give further details on these results.

In contrast, minimax optimal, sparsity-adaptive estimation of the potentially sparse mixture weights has only been studied very recently in [8]. Given an estimator $\widehat{A}$ with $\widehat{A}_{k}\in\Delta_{p}$ for all $k\in[K]$, this work advocates the usage of an MLE-type estimator $\widehat{\alpha}^{(i)}$,

If $\widehat{A}$ were treated as known, and replaced by $A$, the estimator in (3) reduces to $\widehat{\alpha}_{A}^{(i)}$, the MLE of $\alpha^{(i)}$ under the multinomial model (1).

Despite being of MLE-type, the analysis of $\widehat{\alpha}^{(i)}$, for $i\in[n]$, is non-standard, and [8] conducted it under the general regime that we will also adopt in this work. We state it below, to facilitate contrast with the classical set-up.

The General Regime stated above is natural in the topic model context. The fact that $A$ must be estimated from the data is embedded in the definition of the topic model. Also, although one can work with a fixed dictionary size, and fixed number of topics, one can expect both to grow as more documents are added to the corpus. As already mentioned, each document in a corpus will typically cover only a few topics, so $\alpha^{(i)}$ is typically sparse. When $A$ is also sparse, as it is generally the case [10], then $r^{(i)}$ can also be sparse. To see this, take $K=2$ and write $r^{(i)}_{1}=\alpha^{(i)}_{1}A_{11}+\alpha^{(i)}_{2}A_{12}$. If document $i$ covers topic 1 ($\alpha^{(i)}_{1}>0$), but not topic 2 ($\alpha^{(i)}_{2}=0$), and the probability of using word 1, given that topic 1 is covered, is zero ($A_{11}=0$), then $r^{(i)}_{1}=0$.

Under the General Regime above, [8, Theorem 10] shows that, whenever $\alpha^{(i)}$ is sparse, the estimator $\widehat{\alpha}^{(i)}$ is also sparse, with high probability, and derive the sparsity adaptive finite-sample rates for $\|\widehat{\alpha}^{(i)}-\alpha^{(i)}\|_{1}$ summarized in Section 4.1 below. Furthermore, [8] suggests estimating $W(\boldsymbol{\alpha}^{(i)},\boldsymbol{\alpha}^{(j)};d)$ by $W(\widehat{\boldsymbol{\alpha}}^{(i)},\widehat{\boldsymbol{\alpha}}^{(j)};d)$, with

for given estimators $\widehat{A}$, $\widehat{\alpha}^{(i)}$ and $\widehat{\alpha}^{(j)}$ of their population-level counterparts, and $i,j\in[n]$. Upper bounds for these distance estimates are derived [8], and re-stated in a slightly more general form in Theorem 2 of Section 4.2. The work of [8], however, does not study the optimality of their upper bound, and motivates our next contribution.

In Theorem 3 of Section 4.3, we establish a minimax optimal bound for estimating $W(\boldsymbol{\alpha}^{(i)},\boldsymbol{\alpha}^{(j)};d)$, for any pair $i,j\in[n]$, from a collection of $n$ independent multinomial vectors whose collective cell probabilities satisfy the topic model assumption. Concretely, for any $d$ on $\Delta_{p}$ that is bi-Lipschitz equivalent to the TV distance (c.f. Eq. 26), we have

where $\tau=\max\{\|\alpha^{(i)}\|_{0},\|\alpha^{(j)}\|_{0}\}$, and $\|\cdot\|_{0}$ denotes the $\ell_{0}$ norm of a vector, counting the number of its non-zero elements, For simplicity of presentation we assume that $N_{i}=N_{j}=N$, although this is not needed for either theory or practice. The two terms in the above lower bound quantify, respectively, the smallest error of estimating the distance when the mixture components are known, and the smallest error when the mixture weights are known.

We also obtain a slightly sharper lower bound for the estimation of the mixing measures $\boldsymbol{\alpha}^{(i)}$ and $\boldsymbol{\alpha}^{(j)}$ themselves in the Wasserstein distance:

As above, these two terms reflect the inherent difficulty of estimating the mixture weights and mixture components, respectively. As we discuss in Section 4.3, a version of this lower bound with additional logarithmic factors follows directly from the lower bound presented above for the estimation of $W(\boldsymbol{\alpha}^{(i)},\boldsymbol{\alpha}^{(j)};d)$. For completeness, we give the simple argument leading to this improved bound in Appendix D of [7].

The upper bounds obtained in [8] match the above lower bounds, up to possibly sub-optimal logarithmic terms. Together, these results show that the plug-in approach to estimating the Wasserstein distance is nearly optimal, and complement a rich literature on the performance of plug-in estimators for non-smooth functional estimation tasks [41, 16]. Our lower bounds complement those recently obtained for the minimax-optimal estimation of the TV distance for discrete measures [36], but apply to a wider class of distance measures and differ by an additional logarithmic factor; whether this logarithmic factor is necessary is an interesting open question. The results we present are the first nearly tight lower bounds for estimating the class of distances we consider, and our proofs require developing new techniques beyond those used for the total variation case. More details can be found in Section 4.3. The results of this section are valid for any $K,p,n,N$. In particular, the ambient dimension $p$, and the number of mixture components $K$ are allowed to grow with the sample sizes $n,N$.

To the best of our knowledge, despite an active interest in inference for the Wasserstein distance in general [46, 22, 23, 19, 20, 21] and the Wasserstein distance between discrete measures, in particular [57, 59], inference for the Wasserstein between sparse discrete mixing measures, in general, and in particular in topic models, has not been studied.

Our main statistical contribution is the provision of inferential tools for $W(\boldsymbol{\alpha}^{(i)},\boldsymbol{\alpha}^{(j)};d)$, for any pair $i,j\in[n]$, in topic models. To this end:

1. (1)

   In Section 5.2, Theorem 5, we derive a $\sqrt{N}$ distributional limit for an appropriate estimator of the distance.

2. (2)

   In Section 5.3, Theorem 6, we estimate consistently the parameters of the limit, thereby providing theoretically justified, fully data driven, estimates for inference. We also contrast our suggested approach with a number of possible bootstrap schemes, in our simulation study of Section F.4 in [7], revealing the benefits of our method.

Our approach to (1) is given in Section 5. In Proposition 1 we show that, as soon as that error in estimating $A$ can be appropriately controlled in probability, for instance by employing the estimator given in Section 4.1, the desired limit of a distance estimator hinges on inference for the potentially sparse mixture weights $\alpha^{(i)}$. Perhaps surprisingly, this problem has not been studied elsewhere, under the General Regime introduced in Section 2.2 above. Its detailed treatment is of interest in its own right, and constitutes one of our main contributions to the general problem of inference for the Wasserstein distance between mixing measures in topic models. Since our asymptotic distribution results for the distance estimates rest on determining the distributional limit of $K$-dimensional mixture weight estimators, the asymptotic analysis will be conducted for $K$ fixed. However, we do allow the ambient dimension $p$ to depend on the sample sizes $n$ and $N$.

In this section, we begin by reviewing (i) model identifiability and the minimax optimal estimation of the mixing components $A$, and (ii) estimation and finite sample error bounds for the mixture weights $\alpha^{(i)}$ for $i\in[n]$. Necessary notation, used in the remainder of the paper, is also introduced here.

For any estimators $\widehat{\alpha}^{(i)},\widehat{\alpha}^{(j)}\in\Delta_{K}$ and $\widehat{A}$, plug-in estimators $W(\widehat{\boldsymbol{\alpha}}^{(i)},\widehat{\boldsymbol{\alpha}}^{(j)};d)$ of $W(\boldsymbol{\alpha}^{(i)},\boldsymbol{\alpha}^{(j)};d)$ with $\widehat{\boldsymbol{\alpha}}^{(i)}$ and $\widehat{\boldsymbol{\alpha}}^{(j)}$ defined in (4) can be analyzed using the fact that

with $\mathrm{diam}(\mathcal{A})\coloneqq\max_{k,k^{\prime}}d(A_{k},A_{k^{\prime}})$, whenever the following hold

and

for some deterministic sequences $\epsilon_{\alpha,A}$ and $\epsilon_{A}$. The notation $\epsilon_{\alpha,A}$ is mnemonic of the fact that the estimation of mixture weights depends on the estimation of $A$.

Inequality (22) follows by several appropriate applications of the triangle inequality and, for completeness, we give the details in Section B.2.1. We also note that (22) holds for $\widehat{\boldsymbol{\alpha}}^{(i)}$ based on any estimator $\widehat{\alpha}^{(i)}\in\Delta_{K}$.

When $\widehat{\alpha}^{(i)}$ and $\widehat{\alpha}^{(j)}$ are the MLE-type estimators given by (3), [8] used this strategy to bound the left hand side in (22) for any estimator $\widehat{A}$, and then for a minimax-optimal estimator of $A$, when $d$ is either the total variation distance or the Wasserstein distance on $\Delta_{p}$.

We reproduce the result [8, Corollary 13] below, but we state it in the slightly more general framework in which the distance $d$ on $\Delta_{p}$ only needs to satisfy

for some $C_{d}>0$. This holds, for instance, if $d$ is an $\ell_{q}$ norm with $q\geq 1$ on $\mathbb{R}^{p}$ (with $C_{d}=2$), or the Wasserstein distance on $\Delta_{p}$ (with $C_{d}=\mathrm{diam}(\mathcal{Y})$). Recall that $\Theta_{A}$ and $\Theta_{\alpha}(\tau)$ are defined in (12) and (15). Further recall $\overline{\kappa}_{\tau}$ from (19) and $L=n\vee p\vee N$.

In the next section we investigate the optimality of this upper bound by establishing a matching minimax lower bound on estimating $W(\boldsymbol{\alpha}^{(i)},\boldsymbol{\alpha}^{(j)};d)$.

Theorem 3 below establishes the first lower bound on estimating $W(\boldsymbol{\alpha}^{(i)},\boldsymbol{\alpha}^{(j)};d)$, the Wasserstein distance between two mixing measures, under topic models. Our lower bounds are valid for any distance $d$ on $\Delta_{p}$ satisfying

with some $0<c_{d}\leq C_{d}<\infty$.

The lower bound in Theorem 3 consists of two terms that quantify, respectively, the estimation error of estimating the distance when the mixture components, $A_{1},\ldots,A_{K}$, are known, and that when the mixture weights, $\alpha^{(1)},\ldots,\alpha^{(n)}$, are known. Combined with Theorem 2, this lower bound establishes that the simple plug-in procedure can lead to near-minimax optimal estimators of the Wasserstein distance between mixing measures in topic models, up to the $\ell_{1}\to\ell_{1}$ condition numbers of $A$, the factor $C_{d}/c_{d}$ and multiplicative logarithmic factors of $L$.