JIRA issue: [SPARK-1405] Gibbs sampling based Latent Dirichlet Allocation (LDA) for MLlib

[yinxusen]

(This PR is based on a joint work done with @liancheng four months ago.)

Overview

LDA is a classical topic model in machine learning, that provides the ability to extract topics from corpus. Gibbs sampling (GS for short) is a common way to optimize LDA model.

The LDA model consists of four matrices, two 1-dim matrices:

• Document counts
• Topic counts

plus two 2-dim matrices:

• Document-Topic counts
• Topic-Term counts

Implementation details

• An accumulator is used to aggregate all updated values and applies them on the old model computed in the last iteration.
• Chalk is used for term segmentation. Though it is easy to rewrite it with Lucene analyzers, I think MLlib should not take the burden to maintain an implementation of tokenizer.
• SparkContext.wholeTextFiles() is convenient for offline experimentation, while SparkContext.textFile() is better for online applications.
• Document dictionary and term dictionary are broadcasted to translate document names and terms into Int IDs.
• Topic assignment matrix from the last iteration is cached for the current iteration, and then unpersisted to release memory.
• LDA suffers similar stack overflow problem of MLlib ALS (SPARK-1006). To workaround this issue, we checkpoint every a few iterations.

[AmplabJenkins]

Merged build triggered.

[AmplabJenkins]

Merged build started.

[AmplabJenkins]

Merged build finished.

[AmplabJenkins]

Refer to this link for build results: https://amplab.cs.berkeley.edu/jenkins/job/SparkPullRequestBuilder/14318/

[etrain]

I expect that this will be really big - maybe the last two variables should be RDDs - similar to what we do with ALS.

[yinxusen]

That's make sense. I think the docTopicCounts could be sliced easily W.R.T. documents partitions. But for topicTermCounts, it's hard to do slice. I'll find a way to settle it.

[etrain]

Before I get too deep into this review - I want to step back and think about whether we expect the model in this case to be on the order of the size of the data - I think it is, and if so, we may want to consider representing the model as RDD[DocumentTopicFeatures] and RDD[TopicWordFeatures], similar to what we do with ALS. This may change the algorithm substantially.

Separately, maybe it makes sense to have a concrete use case to work with (reuters dataset or something) so that we can evaluate how much memory actually gets used given a reasonably sized corpus.

Perhaps @mengxr or @jegonzal has a strong opinion on this.

[etrain]

Also, speaking of @jegonzal maybe this is a natural first point of integration between MLlib and GraphX - I know the GraphX folks have an implementation of LDA, and maybe this is a chance for us to leverage that work.

[yinxusen]

Yep, I know @jegonzal for his paper Parallel Gibbs Sampling. But I only have the idea of the implementation on GraphLab and not find the impl in GraphX. It's great if I have the chance to talk with Joseph offline.

Besides, I will add a use case for reuters dataset and try to fix the issues put above.

[jegonzal]

I would be happy to talk more about this after the OSDI deadline. As far as storing the model (or more precisely the counts and samples) as an a RDD, I think this really is necessary. The model in this case should be on the order of the size of the data.

Essentially what you want is the ability to join the term topic counts with the document topic counts for each token in a given document. Given these two counts tables (along with the background distribution of topics in the entire corpus) you can compute the new topic assignment.

Here is an implementation of the collapsed Gibbs sampler for LDA using GraphX: amplab/graphx#113

[yinxusen]

Yep, thanks @jegonzal and @etrain , I'll try to fix these issues and look forward to the next step updating and discussion.

[witgo]

We should allow users to customize here. We can add a parameter tokenizer: (String) => Iterable[String] to loadCorpus, and dirStopWords is not required.

[mengxr]

@yinxusen Per discussion on https://issues.apache.org/jira/browse/SPARK-1405, we want to have a GraphX-based implementation and distributed representation of the topic model. Do you mind closing this PR? Thanks for your contribution and @etrain @jegonzal and @witgo for code review!