Spark MLlib LDA 基于GraphX实现原理及源码分析
LDA背景
LDA（隐含狄利克雷分布）是一个主题聚类模型，是当前主题聚类领域最火、最有力的模型之一，它能通过多轮迭代把特征向量集合按主题分类。目前，广泛运用在文本主题聚类中。
LDA的开源实现有很多。目前广泛使用、能够分布式并行处理大规模语料库的有微软的LightLDA，谷歌plda、plda+，sparkLDA等等。下面介绍这3种LDA：
LightLDA依赖于微软自己实现的multiverso参数服务器，服务器底层使用mpi或zeromq发送消息。LDA模型(word-topic矩阵)由参数服务器保存，它为文档训练进程提供参数查询、更新服务。
plda、plda+使用mpi消息通信，将mpi进程分为word、doc俩部分。doc进程训练文档，word进程为doc进程提供模型的查询、更新功能。
spark LDA有两种实现：1.基于gibbs sampling原理和使用GraphX实现的版本（即spark文档上所说的EMLDAOptimizer and DistributedLDAModel），2.基于变分推断原理实现的版本（即spark文档上的OnlineLDAOptimizer and LocalLDAModel）。
LightLDA，plda、plda+，spark LDA比较
论能够处理预料库的规模大小，LihgtLDA要远远好于plda和spark LDA
经过测试，在10个服务器（8核40GB）集群规模下:
LihgtLDA能够处理上亿文档、百万词汇的语料库，能够训练上百万主题数。这样的处理能力使得LihgtLDA能够轻松训练绝大多数语料库。微软号称使用几十机器的集群便能训练Bing搜索引擎爬下数据的十分之一。
相对于LihgtLDA ，plda+能够处理规模小的多，上限是：词汇数目*主题数（模型大小） < 5亿。当语料库规模达到上限后，mpi集群会因内存不够而终止，或因为内存数据频繁切换，迭代速度十分缓慢。虽然plda+对语料库的词汇数目和训练的主题数目很敏感，但对文档的规模并不是很敏感，在词汇数目和主题数目较小的情况下，1000万级别的文档也能够轻松解决。
spark LDA的GraphX版处理规模衡量标准是图的顶点数据，即（文档数 + 词汇数目）*主题数目，上限是 文档数*主题数 < 50亿（由于词汇数目相对于文档数目往往较小，近似等于 文档数*主题数）。当超过这个规模后，spark集群进入假死状态。不停有节点出现OOM，直至任务以失败告终。
变分推断实现的spark LDA瓶颈是 词汇数目*主题数目，这个值也就是我们所说的模型大小，上限约1亿。为什么存在这个瓶颈呢？是因为变分推断的实现过程中，模型使用矩阵本地存储，各个分区计算模型的部分值，然后在driver上将矩阵reduce叠加。当模型过大，driver节点的内存就无法承受各个分区发过来的模型。
收敛速度上，LightLDA要远快于plda、plda+和spark LDA。小规模语料库（30万文档，10万词，1000主题）测试，LightLDA : plda+ : spark LDA(graphx) = 1:4:50
为什么各种LDA的能够处理语料库规模的衡量标准不一样呢？这与它们的实现方式有关，不同的LDA有不同的瓶颈，我们这里单讲spark LDA，其他lda后续介绍。
spark LDA
spark机器学习库MLlib实现了2个版本的LDA，这里分别叫做Spark EM LDA和Spark Online LDA。它们使用同样的数据输入，但是内部的实现和依据的原理完全不同。Spark EM LDA使用GraphX实现的，通过对图的边和顶点数据的操作来训练模型。而Spark Online LDA采用抽样的方式，每次抽取一些文档训练模型，通过多次训练，得到最终模型。在参数估计上，Spark EM LDA使用gibbs采样原理估计模型参数，Spark Online LDA使用贝叶斯变分推断原理估计参数。在模型存储上，Spark EM LDA将训练的主题-词模型存储在GraphX图顶点上，属于分布式存储方式。Spark Online使用矩阵来存储主题-词模型，属于本地模型。通过这些差异，可以看出Spark EM LDA和Spark Online LDA的不同之处，同时他们各自也有各自的瓶颈。Spark EM LDA在训练时shuffle量相当大，严重拖慢速度。而Spark Online LDA使用矩阵存储模型，矩阵规模直接限制训练文档集的主题数和词的数目。另外，Spark EM LDA每轮迭代完毕后更新模型，Spark Online LDA每训练完抽样的文本更新模型，因而Spark Online LDA模型更新更及时，收敛速度更快。
Spark EM LDA之GraphX实现原理
Spark EM LDA基于gibbs采样原理估计参数，凡是基于gibbs采样原理推断参数的LDA训练过程大都如下：
LDA中文档里的每个词都属于一个主题，LDA训练过程的大体思路是，一轮迭代中，为每篇文档里的每一个词重新选择主题，选择的依据是gibbs采样公式，详细原理参见Parameter estimation for text analysis这篇文章。
LDA实现算法的核心是，为每篇文档的每个词重新选取主题。这个过程GraphX做了巧妙的实现，它以文档到词作为边，以词频作为边数据，把语料库构造成图，把对语料库中每篇文档的每个词操作转化为在图中每条边上的操作，而对边RDD处理是GraphX中最常见的的处理方法。
GraphX把 nkm 、 nkt 矩阵存储在文档顶点和词顶点上，把词频信息存储在边上。它把整个文档聚类结果矩阵、模型矩阵和语料库词频矩阵都表达在图结构中，把LDA算法处理过程表达为对边的遍历处理过程。由于基于gibbs采用的LDA可方便的建模成图，又由机器学习的多轮迭代性质，Spark将其简单高效地实现在GraphX之上，形成了Spark MLlib LDA。
Spark EM LDA初始化
Spark LDA的输入数据为词频矩阵RDD[(Long, Vector)]，其存储格式如下表所示：
为了将文档顶点和词顶点统一编号，Spark LDA将文档顶点和词顶点的顶点ID进行了分配。文档顶点ID编号从0递增，词顶点编号从-1递减。上表中词频矩阵转换为下表所示:
Spark LDA根据文档到边的关系生成的GraphX边，如下图所示，边的格式为[(源顶点ID，目的顶点ID，词频)]，如下表所示。
(0, -1, 2.0), (0, -2, 1.0), (0, -3, 3.0), (0, -4, 4.0) …
(1, -1, 3.0), (1, -2, 0.0), (1, -3, 2.0), (1, -4, 5.0) …
…- Spark LDA边构建
M 语料库中文档数目
V 词频矩阵中词的数目
D M*V词频矩阵
1 for document m from [0,M-1]:
2 for word w in document m, w from [-1,-V]:
3 generate an edge (m, w, D[m][w]) as an element of EdgeRDD将预料库中所有文档到词构建成RDD[(Long, Long, Double)]，GraphX进一步在RDD分区中建立索引，进行优化，形成边RDD。
- Spark LDA顶点向量构建
GraphX使用边RDD初始化顶点RDD。
Spark LDA初始完后，语料库被描述为GraphX的图对象，它拥有包括文档和词的顶点RDD和文档指向词的边RDD。顶点RDD拥有一个K维主题分布向量，边拥有词频数据。
Spark LDA迭代
伪代码：
源码：
val sendMsg: EdgeContext[TopicCounts, TokenCount, (Boolean, TopicCounts)] => Unit = (edgeContext) => {
// Compute N_{wj} gamma_{wjk}
val N_wj = edgeContext.attr
// E-STEP: Compute gamma_{wjk} (smoothed topic distributions), scaled by token count
// N_{wj}.
val scaledTopicDistribution: TopicCounts =
computePTopic(edgeContext.srcAttr, edgeContext.dstAttr, N_k, W, eta, alpha) *= N_wj
edgeContext.sendToDst((false, scaledTopicDistribution))
edgeContext.sendToSrc((false, scaledTopicDistribution))
}
// The Boolean is a hack to detect whether we could modify the values in-place.
// TODO: Add zero/seqOp/combOp option to aggregateMessages. (SPARK-5438)
val mergeMsg: ((Boolean, TopicCounts), (Boolean, TopicCounts)) => (Boolean, TopicCounts) =
(m0, m1) => {
val sum =
if (m0._1) {
m0._2 += m1._2
} else if (m1._1) {
m1._2 += m0._2
} else {
m0._2 + m1._2
}
(true, sum)
}
// M-STEP: Aggregation computes new N_{kj}, N_{wk} counts.
val docTopicDistributions: VertexRDD[TopicCounts] =
graph.aggregateMessages[(Boolean, TopicCounts)](sendMsg, mergeMsg)
.mapValues(_._2)源码中sendMsg对应于伪代码中7-8行，mergeMsg对应于伪代码第9行。
伪代码中第2步，计算所有词顶点的向量叠加值WV。Spark对顶点RDD使用filter算子过滤，得到词顶点RDD。再对词顶点RDD的values调用fold进行顶点向量求和。scala代码实现如下：
graph.vertices.filter(isTermVertex).values.fold(BDV.zeros[Double](numTopics))(_ += _)伪代码中第3-9步则是由GraphX的aggregateMasseges方法来实现。第3-8步，属于aggregateMasseges map阶段，它由边三元组edge(srcId, dv, freq, dstId, wv)生成消息msg（k维向量），并将msg发往两端顶点。aggregateMasseges消息产生依据的公式是伪代码中的第4-5步，它跟LDA gibbs实现中第13步，主题选取依据的gibbs取样公式是一样的。
第9步，属于aggregateMasseges reduce阶段，顶点将收到的所有msg叠加。用户传入的消息聚合函数mergeMsg是向量相加。
综上，Spark LDA迭代过程主要分为两个过程，（1）计算所有词顶点数据叠加的值——向量WV；（2）调用aggregateMasseges方法进行消息产生、发送、聚合； （3）顶点将聚合后的消息作为其新的顶点数据。
Spark LDA实现与LDA gibbs实现算法略有差别。LDA gibbs实现算法中第13步为文档每个单词选择主题之前，会取消词当前选择主题，重新选择完主题后更新全局计数器 nkm 、 ntk 、 nm 、 nk 。LDA实现算法中，每个词处理完毕，文档-主题分布 nkm 、主题-词分布 ntk 都会改变，从而影响下一个词的主题选取。因而，LDA gibbs实现算法是一个细粒度的算法。另外，LDA gibbs算法输入是文档分词后的向量，并非词频矩阵。因而，文档中相同的词会先后多次处理。
Spark LDA在迭代过程中，顶点向量和词顶点叠加向量WV保持不变，等效于全局计数器 nkm 、 ntk 、 nm 、 nk 在迭代过程中保持不变，即文档-主题分布、主题-词分布保持不变。map时，文档中所有词都产生主题选取msg（k维向量），reduce会将该msg聚集（叠加）到相应文档顶点和词顶点上。迭代完成后，以聚集后的顶点向量作为图的顶点数据，从而完成统计数据的更新。Spark LDA每轮迭代完成之后全局计数器 nkm 、 ntk 、 nm 、 nk 才做一次更新，更新结果只能影响下一轮迭代各词的msg计算。因而，Spark LDA实现是一个粗粒度地控制主题分布影响的算法。另外，Spark LDA的输入文档是词频矩阵，所以伪代码中第6步需要将归一化的主题分布乘以词频，相当于LDA实现算法对文档中相同词的多次处理。
优缺点分析
Spark LDA使用图顶点存储文档、词的主题统计信息，以图边来存储文档和词的关系，以遍历图边的方式训练模型，完美表达了LDA训练逻辑。然而，GraphX的aggregateMasseges 方法处理边三元组需要从两端顶点拉取顶点数据，生成消息后把需要消息发往两端顶点。这里拉取数据和消息汇聚都会引起大量的shuffle过程。Spark shuffle是指两个分区间的数据移动。shuffle完整过程分为：（1）发送方将待发送的数据写入本地磁盘；（2）数据序列化后经网络传输；（3）接收方接收流流数据，写在本地磁盘；（4）接收方反序列化数据。这里涉及两次序列化、两次读写磁盘，序列化耗CPU，读写磁盘耗时间。大量的shuffle大大地拖累了Spark LDA的迭代速度。因而，GraphX的性能并不高，不及使用MPI进行消息通信的LightLDA和plda。可以说，GraphX为LDA在Spark上实现提供一个完美的结构，而性能却不敢恭维！！但是，spark由于其强大的普适性，为了减少数据多平台跨越的烦恼，在可接受范围内，使用Spark训练语料库还是可行的。
作者介绍
唐黎哲，国防科学技术大学 并行与分布式计算国家重点实验室(PDL)硕士，从事spark、图计算、文本分析研究，欢迎交流，请多指教。
邮箱：[email protected]