Evaluate topic models over corpora and numbers of topics

Fit one structural topic model per combination of corpus (e.g. vocabularies built with different filters), number of topics K and run (seed), in parallel, and compute the metrics asked for. This is the general version of stm::searchK(): with one corpus, one run and metrics = c("heldout", "residual", "coherence", "exclusivity"), it returns the same values as searchK() (tested).

Available metrics (see the documentation of each function):

"heldout": held-out likelihood (Wallach et al. 2009): part of the words of some documents is hidden during fitting, and the model's expected log-probability of these words is computed (stm::eval.heldout()). Higher is better. Needs heldout = TRUE;
"residual": dispersion of the multinomial residuals (Taddy 2012, stm::checkResiduals()): above 1, the model leaves too much variance, which suggests K is too small;
"bound": final variational lower bound (and lbound, plus \(\log K!\)) and number of EM iterations;
"coherence": topic_coherence() (UMass and NPMI);
"exclusivity": topic_exclusivity();
"diversity": topic_diversity();
"similarity": mean and maximum cosine similarity between topics (topic_similarity(); the mean is Cao et al.'s 2009 criterion);
"ldatuning": ldatuning_metrics() (Arun et al. 2010, Deveaud et al. 2014);
"stability": topic_stability() across the n_runs runs (needs n_runs >= 2 and a random initialisation).

Usage

evaluate_topic_models(
  corpora,
  K,
  metrics = c("heldout", "residual", "bound", "coherence", "exclusivity", "diversity",
    "similarity", "ldatuning"),
  ...,
  n_runs = 1L,
  heldout = TRUE,
  n_heldout = NULL,
  heldout_proportion = 0.5,
  heldout_seed = seed,
  n_terms = 10L,
  frexw = 0.7,
  n_workers = 1L,
  output_dir = NULL,
  save_models = FALSE,
  seed = 1L
)

Arguments

corpora: A named list of STM inputs, each a list with documents, vocab and, if the formula needs it, meta (e.g. from tokens_to_stm()). A single STM input is accepted and named "corpus".
K: Integer vector, numbers of topics.
metrics: Character vector of metrics (see Description). Default: all but stability (which needs several runs).
...: Further arguments of stm::stm(): prevalence, content, max.em.its, emtol, init.type... The data argument is taken from each corpus's meta; do not give it here. Use the settings of the final model.
n_runs: Integer, number of runs per corpus and K, with seeds seed, seed + 1, ... Needed for stability. Spectral initialisation (stm's default) is deterministic: use init.type = "LDA" or "Random" for several runs.
heldout: Logical: fit on documents with held-out words (needed for the held-out likelihood, as searchK() does) or on the full documents. With TRUE, every metric is computed on the model fitted with held-out words.
n_heldout, heldout_proportion: Number of documents with held-out words (default: 10% of the documents) and share of their words held out (stm::make.heldout()).
heldout_seed: Integer seed of the held-out sample (one sample per corpus, shared by every K and run, so that likelihoods are comparable across K). Default: seed. The sample must be the same from one call to the next for output_dir to resume.
n_terms, frexw: Top words used by coherence, exclusivity and stability, and FREX weight of exclusivity (as searchK(): 10 and 0.7).
n_workers: Integer, number of models fitted at the same time. Each worker holds a copy of the corpora and fits one model: peak memory is about n_workers times the memory of one fit. Set it from the memory available, not from the number of cores.
output_dir: Optional directory. When given, each fitted corpus x K x run is saved there as soon as it is done (<corpus>_K<K>_run<run>.rds, with its metrics), and combinations already saved with the same corpus and settings are not refitted: an interrupted evaluation resumes where it stopped. NULL (default) writes nothing.
save_models: Logical: also save the fitted models in output_dir (e.g. to prepare intrusion tests, make_word_intrusion(), without refitting). Models can be large.
seed: Integer, base seed of the stm fits (seed + run - 1) and of the workers' random number streams.

Value

A list of class topic_model_evaluation with:

models: one row per corpus, K and run: corpus, K, run, the metrics, n_documents, vocabulary_size, minutes;
topics: one row per corpus, K, run and topic, with the per-topic metrics (coherence, exclusivity, prevalence, top terms);
stability: one row per corpus, K and topic of the reference run, when "stability" is asked;
settings: the arguments used. Failed fits are absent, with a warning giving the error.

Details

Choosing K is a judgement, not an optimisation. The metrics pull in different directions: coherence favours few broad topics, exclusivity and held-out likelihood often favour many. The usual practice (Roberts, Stewart and Tingley 2019; Grimmer, Roberts and Stewart 2022) is to narrow the candidates with these metrics (plot_model_metrics(), plot_k_frontier()), then to read the topics of a few candidates and validate them with humans (make_word_intrusion()).

Comparing corpora. Metrics computed on different vocabularies are not directly comparable: held-out likelihood is a per-word log-probability over a vocabulary of a given size; coherence depends on document frequencies of the kept words. Compare K within a corpus, and corpora on the shape of their curves and on human validation (Denny and Spirling 2018 show how much preprocessing choices can change the results).

Parallelism and memory. Fits run on n_workers background R processes (mirai). The corpora are sent once to each worker, then each task only carries a corpus name, K and run; tasks are queued, largest K first. A failed fit does not stop the others. Workers load mintR (the development version when the package is loaded with devtools::load_all()).

Content covariates. With content =, each topic has one word distribution per level of the covariate. Coherence, exclusivity, diversity, similarity and the ldatuning criteria are then computed on the marginal distribution of each topic (marginal_topic_model()). Fits with a content covariate are much slower (50 minutes instead of 1.5 for K = 15 on 2,500 speeches and 15,000 terms): reduce the vocabulary or max.em.its for the grid.

References

Denny, M. J. and Spirling, A. (2018). Text preprocessing for unsupervised learning: Why it matters, when it misleads, and what to do about it. Political Analysis, 26(2), 168-189. doi:10.1017/pan.2017.44

Grimmer, J., Roberts, M. E. and Stewart, B. M. (2022). Text as Data: A New Framework for Machine Learning and the Social Sciences. Princeton University Press.

Roberts, M. E., Stewart, B. M. and Tingley, D. (2019). stm: An R package for structural topic models. Journal of Statistical Software, 91(2), 1-40. doi:10.18637/jss.v091.i02

Taddy, M. (2012). On estimation and selection for topic models. Proceedings of AISTATS 2012, 1184-1193.

Wallach, H. M., Murray, I., Salakhutdinov, R. and Mimno, D. (2009). Evaluation methods for topic models. Proceedings of ICML 2009, 1105-1112.

Examples

if (FALSE) { # \dontrun{
evaluation <- evaluate_topic_models(
  corpora = list(strict = stm_strict, loose = stm_loose),
  K = seq(20, 80, by = 10),
  prevalence = ~ s(year), max.em.its = 150,
  heldout_seed = 2026, n_workers = 3, output_dir = "data/evaluation"
)
plot_model_metrics(evaluation)
plot_k_frontier(evaluation)
} # }