[论文解读] Fast MLE Computation for the Dirichlet Multinomial

Max Sklar|arXiv (Cornell University)|May 1, 2014

Bayesian Methods and Mixture Models参考文献 13被引用 12

一句话总结

本文提出一种新颖的单遍算法，通过预计算 digamma 和 trigamma 函数项，利用牛顿-拉夫森优化实现 Dirichlet-mulitnomial 参数的快速最大似然估计（MLE）。通过将预计算与迭代优化解耦，显著降低运行时间，尤其在稀疏数据下效果显著，实现与数据集大小无关的常数时间牛顿-拉夫森步骤，实测速度提升最高达10倍，优于标准实现。

ABSTRACT

Given a collection of categorical data, we want to find the parameters of a Dirichlet distribution which maximizes the likelihood of that data. Newton's method is typically used for this purpose but current implementations require reading through the entire dataset on each iteration. In this paper, we propose a modification which requires only a single pass through the dataset and substantially decreases running time. Furthermore we analyze both theoretically and empirically the performance of the proposed algorithm, and provide an open source implementation.

研究动机与目标

• 加速大规模数据应用中 Dirichlet-mulitnomial 分布的 MLE 计算。
• 解决标准牛顿-拉夫森方法每轮迭代需完整扫描数据集导致的低效问题。
• 在高维与稀疏场景下，实现可扩展且可并行化的 Dirichlet 参数估计。
• 支持贝叶斯模型、主题模型（如 LDA）和混合模型中的高效推断。
• 提供实用、开源的实现，适用于生产环境与研究用途。

提出的方法

• 提出两阶段算法：首先利用 digamma 和 trigamma 函数预计算行级统计量（U 和 v），随后进行牛顿-拉夫森优化。
• 利用 digamma 和 trigamma 函数的数学恒等式，预计算原本每轮迭代需重新计算的项。
• 将预计算步骤与牛顿-拉夫森迭代解耦，实现行的独立处理，支持并行化。
• 使用预计算的 U 和 v 对全数据集应用牛顿-拉夫森方法，收敛性与行数 M 无关。
• 对幂律分布数据采用混合方法，将高计数行与低计数行分离处理，以优化性能。
• 通过维护 U 和 v 作为数据的运行汇总，支持增量更新与提前停止。

实验结果

研究问题

• RQ1 通过预计算充分统计量，能否加速 Dirichlet-mulitnomial 分布的 MLE 计算？
• RQ2 将预计算与牛顿-拉夫森分离，是否能显著降低整体运行时间，尤其在稀疏或大规模数据下？
• RQ3 随着类别数 K 和多项分布样本数 M 增加，该方法的可扩展性如何？
• RQ4 该算法能否高效并行化或适配流式数据处理？
• RQ5 在真实场景中，与 Minka 或 Wallach 的定点迭代方法相比，性能表现如何？

主要发现

• 所提算法在标准牛顿-拉夫森实现上实现最高达10倍的速度提升，尤其在 M 较小而 N 较大时效果显著。
• 随着 M 增大，牛顿-拉夫森步骤运行时间保持恒定，因为不再需要在每轮迭代中读取完整数据集。
• 预计算高度可并行化，可提前停止或增量扩展，支持流式与分布式处理。
• 该方法保持数值精度，收敛至与其它基于牛顿-拉夫森的方法相同的 MLE 解，且初始化条件一致。
• 当 K 极高时（如 2048 维），性能下降，可能由于内存分配与数据结构开销。
• 将 Minka 方法用于高计数行、本方法用于低计数行的混合策略，性能优于单独使用任一方法。