Arora 2012 - Recover Algorithms RecoverKL and RecoverL2

Summary

With anchors $S$ known, every row of the row-normalised co-occurrence matrix $\bar{Q}$ lies in the convex hull of the anchor rows, and the mixing weights equal $p(z_1\mid w_1=i)$. Solving a K-dimensional simplex-constrained regression per word, then applying Bayes’ rule, yields $A$ without matrix inversion.

Main Content

Input. $Q\in {\mathbb{R}}^{V\times V}$ of word-word co-occurrence, normalised to sum 1: in the population limit $Q_{ij}=p(w_1=i,w_2=j)$ and $Q=E[AW{W}^T{A}^T]=AR{A}^T$. Row-normalised ${\bar{Q}}_{ij}=p(w_2=j\mid w_1=i)$. Construction (Appendix D.1): for each document with counts $H_d$ and length $n_d$, use $\frac{1}{n_d(n_d-1)}(H_d{H}_d^T-\text{Diag}(H_d))$, which has expectation $A{W}_d{W}_d^T{A}^T$ (removes the diagonal word-with-itself terms), averaged over documents.

Convex-hull structure

For an anchor word $s_k$: ${\bar{Q}}_{s_k,j}=p(w_2=j\mid z_1=k)$, since $p(z_1=k\mid w_1=s_k)=1$ and $w_2\perp w_1\mid z_1$ in an admixture. For any word $i$: ${\bar{Q}}_{i,j}={\sum }_k{C}_{ik}{\bar{Q}}_{s_k,j}$ with $C_{ik}=p(z_1=k\mid w_1=i)\ge 0$, ${\sum }_k{C}_{ik}=1$. Hence every row of $\bar{Q}$ lies in the convex hull of the anchor rows.

Original Recover (Alg. 2) and RecoverKL (Alg. 3)

Alg. 2 (Arora et al. 2012b): permute $Q$ so anchors come first, compute ${\vec{p}}_S=Q_S\vec{1}$, solve $Q_{S,S}\vec{z}={\vec{p}}_S$, then $A^T=(Q_{S,S}\text{Diag}(\vec{z}){)}^{-1}{Q}_S^T$ and $R=\text{Diag}(\vec{z})Q_{S,S}\text{Diag}(\vec{z})$ - matrix inversion; can return negative values. RecoverKL: for each word $i$, $C_{i\cdot }=\arg {\min }_{C_i}{D}_{KL}({\bar{Q}}_i\Vert {\sum }_{k\in S}{C}_{ik}{\bar{Q}}_{s_k})$ subject to ${\sum }_k{C}_{ik}=1$, $C_{ik}\ge 0$ (tolerance $ϵ$); then $A^{\prime }=\text{diag}({\vec{p}}_w)C$ with ${\vec{p}}_w=Q\vec{1}$ and columns of $A^{\prime }$ normalised to give $A$ (Bayes: $p(w=i\mid z=k)\propto p(z=k\mid w=i)p(w=i)$); $R=A^{†}Q{A}^{†T}$. RecoverL2: same with squared loss; objective $\Vert Q_i-C_i^T{Q}_S{\Vert }^2=\Vert Q_i{\Vert }^2-2C_i^T(Q_S{Q}_i^T)+C_i^T(Q_S{Q}_S^T)C_i$ where $Q_S{Q}_S^T$ is $K\times K$ and computed once, so the optimisation is independent of $V$.

Details: the per-word problems are independent (parallel) and solved by exponentiated gradient (Alg. 5, Kivinen and Warmuth 1995) with a KKT-based convergence test (stationarity, primal/dual feasibility, complementary slackness ${\lambda }^{T}x<ϵ$, equivalent to a small primal-dual gap). RecoverL2 runtime $O(K{V}^2+K^{2}VT)$ with $T\approx 100$-$1000$ iterations. KL form is maximum likelihood on the co-occurrence counts but does not force the admixture constraint. For LDA, $p(z)={\sum }_{i}p(z\mid w=i)p(w=i)$ gives the Dirichlet hyperparameters up to scale; the scale is chosen by grid search on training likelihood (more reliable than reading it from $R$).

Examples

Population-level recovery (runnable)

Synthetic $V=40$, $K=3$ with exact anchors; exact $Q=AR{A}^T$; anchors via FastAnchorWords; recovery with simplex-constrained least squares gives L1 error 0 (to numerical precision) per topic.

import numpy as np
from scipy.optimize import nnls
rng = np.random.default_rng(0)
V, K = 40, 3
A = rng.dirichlet(np.ones(V) * 0.2, K).T
A[:K, :] = 0; A[np.arange(K), np.arange(K)] = 0.05; A /= A.sum(0)   # words 0..K-1 are anchors
R = 0.6 * np.eye(K) + 0.4 / K
Q = A @ R @ A.T; Q /= Q.sum(); Qb = Q / Q.sum(1, keepdims=True)
 
def fast_anchor_words(Qb, K):
    S = [int(np.argmax((Qb**2).sum(1)))]
    for _ in range(K - 1):
        B, _ = np.linalg.qr(Qb[S].T)
        resid = Qb.T - B @ (B.T @ Qb.T)
        S.append(int(np.argmax((resid**2).sum(0))))
    return S
 
def recover_l2(Q, Qb, S, lam=1e3):
    pw = Q.sum(1); T = Qb[S].T; K = len(S)
    C = np.array([nnls(np.vstack([T, lam*np.ones(K)]), np.append(Qb[i], lam))[0] for i in range(Q.shape[0])])
    C /= C.sum(1, keepdims=True)
    A_hat = C * pw[:, None]
    return A_hat / A_hat.sum(0)
 
S = fast_anchor_words(Qb, K)
print(sorted(S), np.abs(recover_l2(Q, Qb, S) - A[:, S]).sum(0).round(3))   # anchors found; ~0 error

Connections

Why it works: Arora 2012 - Separability and Anchor Words. Anchor selection: Arora 2012 - FastAnchorWords Algorithm and Guarantees. Compare to likelihood-based fitting: Blei 2003 - Variational Inference and Parameter Estimation.

See Also

Arora 2012 - Practical Provable Topic Modeling - Overview