Non-negative matrix factorization algorithms generally improve topic model fits

In plain words

A topic model splits a table of counts, such as words per document or genes per cell, into a few "topics" that each row mixes in different amounts. The standard way to fit one is an algorithm called EM. It can be very slow, and it sometimes stalls far away from a good answer.

The paper notes that a topic model is closely related to a simpler problem, non-negative matrix factorisation (NMF). NMF has no "shares must add up to one" rule, so its optimisers are easier to run. The fit is done in the simpler form and then translated back into topics.

The authors show that a classic NMF update rule is, in different wording, the same EM algorithm used for topic models. They also find that a coordinate-descent optimiser for NMF, with an extra acceleration step, usually reaches better fits faster than EM. Their code is the R package fastTopics.

TL;DR

• Maximum-likelihood fits of the multinomial topic model match those of Poisson NMF, once the parameters are mapped between the two models §2 eq. 1 eq. 6.
• Because Poisson NMF drops the sum-to-one constraints, NMF optimisers can be used directly §2 eq. 5 eq. 7.
• The NMF multiplicative updates are EM for the topic model, up to the order of operations §3 tab. 1.
• In the authors' runs on four count data sets, coordinate descent with extrapolation was usually at least as fast as the other variants and often reached the best fit §4 fig. 2. These results are author-reported.

Problem

The aim is the maximum-likelihood estimate (MLE) of a multinomial topic model for an n × m count matrix X, with K ≥ 2 topics §1 §2. The model has row-stochastic loadings L and column-stochastic factors F eq. 5 eq. 6. Fitting it means solving a linearly constrained problem eq. 7.

The EM algorithm for this problem can converge very slowly §1. In the MCF-7 example, 200 EM updates left the estimate more than 150,000 log-likelihood units from the MLE, while 200 coordinate-descent updates were within 0.079 §4 fig. 1. EM eventually reached the same MLE after many more updates §4. The authors argue that such slow progress can change the topics, not merely delay convergence §4 §5.

Assumptions: the paper treats point estimates (MLE, and MAP with Dirichlet priors on F) and does not address full posterior inference §5.

Method

Models and the change of variables

Poisson NMF. Counts follow x_ij ~ Pois(λ_ij) with λ_ij = Σ_k h_ik w_jk, where H (n × K) and W (m × K) are non-negative eq. 1. Its MLE solves a bound-constrained problem eq. 2 with objective ℓ eq. 3, which uses φ eq. 4.

Topic model. Counts follow a multinomial with totals t_i = Σ_j x_ij and probabilities π_ij = Σ_k l_ik f_jk eq. 6. Each row of L and each column of F sums to one eq. 5. The MLE is a linearly constrained problem with objective φ eq. 7.

Mapping (PNMF-to-MTM, Definition in §2). Given H and W:

1. Normalise the columns of W. Keep the column totals u and the normalised columns as F.
2. Form H·diag(u). Normalise its rows. Keep the row totals s and the normalised rows as L.

The inverse is W = F·diag(u) and H = diag(s)·L·diag(u)⁻¹ §2. The map is a one-to-one change of variables when every row of H and every column of W has a positive entry §2.

Lemma §2 eq. 8. The Poisson NMF likelihood equals the topic-model likelihood times Π_i Pois(t_i; s_i). The proof uses the identity that a product of Poissons equals a multinomial times a Poisson of the total eq. 9.

Corollary §2. Poisson NMF MLEs map to topic-model MLEs under the mapping above. Conversely, topic-model MLEs with s = t and any positive u map back to Poisson NMF MLEs. The inverse step is unique only when every t_i > 0 §2 (proof).

Non-identifiability. H and W are unique only up to rescaling column k of H by a and column k of W by 1/a §2 remark. Raw H and W are therefore not comparable across topics; the L and F from the mapping are §2.

Alternating Poisson regression

Fitting Poisson NMF by block updates is Algorithm 1 alg. 1. Each row of H is the MLE of an additive Poisson regression with the rows of W as design matrix, and then each row of W is fitted with the new H fixed alg. 1 §3. The row subproblems are independent and can run in parallel §3.

The additive Poisson regression model is y_i ~ Pois(μ_i) with μ_i = Σ_k a_ik b_k and b ≥ 0 §3. Two solvers for it are discussed:

• EM §3. The E-step computes z̄_ik = y_i a_ik b_k / μ_i. The M-step sets b_k to Σ_i z̄_ik divided by Σ_i a_ik. Combining the steps gives b_k ← b_k · Σ_i a_ik y_i / μ_i ÷ Σ_i a_ik. Applied to the rows of H and W, this recovers the Poisson NMF multiplicative updates §3. Applying the mapping to those updates gives EM for the topic model §3.
• Coordinate descent (CD) §3. Each b_k is updated by b_k ← max{0, b_k − α_k g_k / q_k}, where g_k = Σ_i a_ik (1 − y_i / μ_i) and q_k = Σ_i y_i a_ik² / μ_i². The variants CCD and SCD take the full step α_k = 1 and use no line search, so the objective need not decrease at each step §3.
• Algorithm map tab. 1. EM for Poisson regression corresponds to multiplicative updates for Poisson NMF and to EM for the topic model. CD corresponds to CCD and SCD for Poisson NMF, and has no counterpart for the topic model because of the sum-to-one constraints §3 tab. 1.

The paper's "EM" and "CD" are the two solvers inserted into Algorithm 1 §3 §4. Within each outer iteration, the updates are applied row by row (stepwise), not as a full E-step followed by a full M-step §3.

Extrapolation

The extrapolated update is H^ext = P₊[H^new + β(H^new − H^(t−1))], and likewise for W. Here P₊ projects onto the non-negative orthant, and β ∈ [0, 1]. Setting β = 0 gives the plain update app. B. The scheme follows Ang and Gillis (2019) app. B. The paper activates it at iteration 50 of CD on MCF-7 §4 and uses it with CD and EM in the main comparisons §4.

Pseudocode

Algorithm 1 alg. 1, with the subproblem solver left abstract:

# Alternating Poisson regression for Poisson NMF  [alg. 1]
# Input:  counts (n, m); initial H (n, K), W (m, K)
# Output: H (n, K), W (m, K), both non-negative

def alternating_poisson_regression(counts, H, W, num_outer):
    # Repeat the outer loop; the paper uses a fixed number of updates  [§4]
    for t in range(num_outer):
        # Refit every row of H with W fixed  [alg. 1]
        for i in range(n):
            H[i] = fit_pois_reg(W, counts[i])     # (K,) from (m, K)
        # Refit every row of W with the new H fixed  [alg. 1]
        for j in range(m):
            W[j] = fit_pois_reg(H, counts[:, j])  # (K,) from (n, K)
    return H, W

One EM sweep for the subproblem, as in the multiplicative update §3:

# One EM update for additive Poisson regression  [§3]
# Input:  y (n,), A (n, K), b (K,) with b > 0
# Output: b (K,)

def em_step(y, A, b):
    # Fitted means  mu_i = sum_k a_ik * b_k  [§3]
    mu = A @ b                                 # (n,)
    # Expected counts  z_ik = y_i * a_ik * b_k / mu_i  [§3]
    z = y[:, None] * A * b[None, :] / mu[:, None]   # (n, K)
    # Update each b_k to expected counts over the column sum of A  [§3]
    b = z.sum(axis=0) / A.sum(axis=0)          # (K,)
    return b

One coordinate-descent sweep, without line search §3:

# One coordinate-descent sweep for additive Poisson regression  [§3]
# Input:  y (n,), A (n, K), b (K,) with b >= 0
# Output: b (K,), non-negative

def cd_sweep(y, A, b):
    for k in range(K):
        # Fitted means at the current b  [§3]
        mu = A @ b                             # (n,)
        # First and second derivatives of the negative log-likelihood  [§3]
        g = (A[:, k] * (1 - y / mu)).sum()     # scalar
        q = (y * A[:, k] ** 2 / mu ** 2).sum() # scalar
        # Full Newton step (alpha = 1), clipped at zero  [§3]
        b[k] = max(0.0, b[k] - g / q)
    return b

Extrapolation of the non-negative estimates app. B:

# Extrapolated update  [app. B]
# Input:  new (n, K), previous (n, K), beta in [0, 1]
# Output: extrapolated estimate (n, K), non-negative

def extrapolate(new, previous, beta):
    # Move beyond the new estimate by beta times its last change  [app. B]
    candidate = new + beta * (new - previous)   # (n, K)
    # Project onto the non-negative orthant  [app. B]
    return np.maximum(candidate, 0.0)           # (n, K)

Recovering topic-model parameters from a Poisson NMF fit §2:

# Map Poisson NMF parameters to a topic model  [§2] [Definition]
# Input:  H (n, K), W (m, K)
# Output: L (n, K) rows sum to 1; F (m, K) columns sum to 1; s (n,)

def pnmf_to_topic_model(H, W):
    # Column-normalise W, keeping the column totals u  [§2]
    u = W.sum(axis=0)                           # (K,)
    F = W / u                                   # (m, K)
    # Row-normalise H diag(u), keeping the row totals s  [§2]
    scaled = H * u                              # (n, K)
    s = scaled.sum(axis=1)                      # (n,)
    L = scaled / s[:, None]                     # (n, K)
    return L, F, s

Hyperparameters

Values marked "paper" are from the text or appendix. "Code" entries are read from the released code and are not stated in the paper.

Data sets tab. 2 app. C. Counts are sparse except MCF-7.

MCF-7 is 41 samples by 16,773 genes, stored dense §4 app. C.

Results

Author-reported. Setup: four count data sets and MCF-7; R implementation; K from 2 to 12 §4 app. C.

The paper gives no numeric table of speed-ups across the four data sets; those comparisons are plotted only §4 fig. 2 fig. 5–fig. 12. It reports the three patterns observed: EM reaches a good solution quickly; the initial EM iterations are insufficient; and EM progresses slowly in a region far from the MLE §4.

Paper vs code

Searched: the paper's own links, github.com/stephenslab/fastTopics (the package) and github.com/stephenslab/fastTopics-experiments. No other search was needed. The stephenslab organisation belongs to the authors' lab (Matthew Stephens is an author). Both were cloned shallowly. The package is at commit 5cf9a448, DESCRIPTION version 0.7-46. The experiments are at commit 2c1efab0. The paper's main numbers used version 0.5-24 app. C, so this card is not checked against results from that version.

Code licence: BSD-2-Clause, LICENSE file in the package repository. Copyright: Peter Carbonetto and Matthew Stephens.

Discrepancies found:

• Extrapolation is off by default. The package default is extrapolate = FALSE fit_poisson_nmf.R. The experiments turn it on explicitly newsgroups.R.
• Extrapolation has a reset. Every extrapolate.reset-th iteration (default 20) uses the plain update fit_poisson_nmf.R fit_poisson_nmf.R. The paper does not mention this.
• Extrapolation schedule. On success β grows by 1.1 and is capped at a value betamax, which grows by 1.05 up to 0.99. On failure β is multiplied by 0.75 and the scheme restarts fit_poisson_nmf.R fit_poisson_nmf.R. The paper gives only β ∈ [0, 1] and says the step size is adapted to performance app. B §4.
• Extrapolated base point. The paper defines H^new as a regression on W^(t−1). The code computes the loading update from the extrapolated factors fit$Fy fit_poisson_nmf.R; the factor update then uses fit$Ly fit_poisson_nmf.R. The two agree only when β = 0.
• Floor. The paper sets estimates below 1e-15 to 1e-15 app. B. The code default floor is 1e-10 fit_poisson_nmf.R.
• Epsilon. The code adds eps = 1e-8 to the loss and to the CD denominators fit_poisson_nmf.R poismix.cpp poismix.cpp. The paper does not mention it. The SCD wrapper's own default is 1e-16, but the fitting path passes control$eps scd.R fit_poisson_nmf.R.
• Which CD. The paper's "CD" is sequential coordinate descent (SCD, from the NNLM package), which is the method = "scd" path app. C fit_poisson_nmf.R. CCD is also in the code (method = "ccd") but the experiments use SCD newsgroups.R.
• Initialisation. The paper says Topic-SCORE initialisation is commented out in the appendix app. C. The package default is init.method = "topicscore" fit_poisson_nmf.R. The prefit script does not set init.method prefit_poisson_nmf.R. The paper does not state the initialisation for the main comparisons.
• Newsgroups preprocessing. The paper describes only the document filter for NeurIPS and newsgroups app. C. The preparation script filters words, not documents prepare_newsgroups.Rmd. The example script newsgroups.R subsamples 4,000 documents and drops words in fewer than 10 documents newsgroups.R newsgroups.R. That script is an example and not the reported run.
• Stopping. The code's convergence tests are disabled by default fit_poisson_nmf.R. The paper uses fixed iteration counts and reports the KKT residual and loss change as diagnostics app. B.
• Consistent. The rescaling rule app. B matches the code misc.R. The topic mapping matches §2 poisson2multinom.R. The 1,000-update prefit matches §4 prefit_poisson_nmf.R. The 4-update inner loop matches app. B.

Pitfalls

• Zero totals break the mapping. The inverse step has a unique solution only when every document total t_i is positive §2. Drop empty rows first. Also avoid zero column sums of W, since the mapping divides by u §2.
• Compare topics only after mapping. Raw H and W are determined only up to column scaling §2 remark. Compare L and F, which the mapping gives.
• Division by the fitted mean. EM and CD divide by μ_i (equal to λ) and by q_k. Near-zero means give NaN or huge steps. The paper sets a floor app. B; the code adds eps fit_poisson_nmf.R. The two values differ, so results near the floor may differ.
• CD may fail to decrease the objective. Full Newton steps have no line search, so single steps can increase the loss far from a solution. The paper runs one EM update before each CD sweep to reduce failures app. B scd.R.
• Extrapolation adds cost. Each extrapolated iteration evaluates the loss again at the extrapolated point fit_poisson_nmf.R. The paper says the overhead is minimal §3.
• Loss cost on sparse data. With N nonzeros the loss costs O((N + n + m)K), against O(nmK) if dense. The paper relies on this for data sets with 0.2% to 9.3% nonzeros app. B tab. 2.
• Incomplete subproblem solves. Each outer step stops after at most four inner updates, warm-started from the previous estimate app. B. The outer iterates are therefore not exact MLEs of the subproblems, and the KKT residual need not fall monotonically app. D fig. 6.
• Rescaling changes H and W. Rescaling leaves λ and the loss unchanged app. B. Do not monitor raw H or W to judge convergence.
• "MLE" is a reference point. Where the true MLE was not reached, the paper uses the best estimate found among runs as reference §4. Distances to "the MLE" are relative to that reference.
• Mixed evidence on EM. The authors separate benign slow EM progress from slow progress far from an MLE. Which case applies can only be seen by comparing algorithms §4 §5.
• Baselines not shown. DAAREM, quasi-Newton EM acceleration and maptpx are described as giving little or no benefit, but the results are "not shown" §5. The maptpx comparison is commented out in the appendix app. D.
• Hardware and version drift. Timings are from R 3.5.1 on Xeon E5-2680v4 with 8 CPUs app. C. The code has since moved to version 0.7-46. Per-iteration time is described as "very similar" across algorithms §4; this card does not verify it.
• MAP extension not checked. The Dirichlet and gamma priors (pseudo-documents and penalties) are derived in the appendix §2 remark app. A. This card did not check their implementation in the code.

Lineage

• Builds on: Lee and Seung (2001) multiplicative updates for Poisson NMF §1 §3; Hofmann (1999) EM for the topic model §1; Cemgil (2009) on the EM equivalence §1; Ang and Gillis (2019) extrapolation app. B; Hien et al. (2020) comparison of Poisson NMF algorithms §3; SCD from Lin and Boutros (NNLM) app. C.
• Compared against: EM for the topic model; plain and extrapolated CD and EM variants §4; DAAREM and quasi-Newton EM acceleration (results not shown) §5. A maptpx comparison appears only as a commented-out figure in the appendix app. D.
• Superseded by: not known from content in this wiki.