[论文解读] Augment and Reduce: Stochastic Inference for Large Categorical Distributions

Francisco J. R. Ruiz, Michalis K. Titsias+2 more|arXiv (Cornell University)|Feb 12, 2018

Bayesian Methods and Mixture Models参考文献 41被引用 5

一句话总结

本文提出了一种通用的随机变分推断方法——增强与缩减（a&r），用于大规模分类分布，通过引入隐变量并优化边际似然的下界，显著降低了计算成本。与 OVE 和精确 softmax 等现有方法相比，该方法在最大达 13,000 个类别的大规模分类任务中，实现了更紧的下界和更优的预测性能。

ABSTRACT

Categorical distributions are ubiquitous in machine learning, e.g., in classification, language models, and recommendation systems. However, when the number of possible outcomes is very large, using categorical distributions becomes computationally expensive, as the complexity scales linearly with the number of outcomes. To address this problem, we propose augment and reduce (A&R), a method to alleviate the computational complexity. A&R uses two ideas: latent variable augmentation and stochastic variational inference. It maximizes a lower bound on the marginal likelihood of the data. Unlike existing methods which are specific to softmax, A&R is more general and is amenable to other categorical models, such as multinomial probit. On several large-scale classification problems, we show that A&R provides a tighter bound on the marginal likelihood and has better predictive performance than existing approaches.

研究动机与目标

• 解决大规模分类分布拟合带来的高计算成本问题，其计算复杂度随类别数 K 线性增长。
• 开发一种通用方法，可扩展至 softmax 之外的模型，包括多项式 probit 和逻辑回归模型。
• 提供边际似然的紧致下界，使其可集成至随机变分 EM 及其他优化框架中。
• 在保持低计算开销的前提下，提升大规模分类任务中的预测性能。
• 实现对输出空间极大规模模型的可扩展训练，例如具有大量类别的图像分类任务或包含海量物品的推荐系统。

提出的方法

• 引入隐变量 ε，将分类分布重参数化为 p(y|ψ) = ∫ p(y,ε|ψ) dε，从而支持变分推断。
• 采用随机变分推断以最大化边际似然的下界，避免直接计算配分函数。
• 应用两步推断流程：局部步骤用于优化 ε 上的变分近似，全局步骤用于更新 ψ。
• 在局部步骤中复用指数函数计算，保持高效性，使计算成本接近 OVE（一对多估计）水平。
• 为 softmax 和多项式 probit 模型分别制定证据下界（ELBO），其中后者使用重要性采样近似积分。
• 将该方法嵌入随机变分 EM 框架中，以支持大规模数据集上的可扩展训练。

实验结果

研究问题

• RQ1 能否开发一种通用方法，将大规模分类分布的推断能力扩展至 softmax 模型之外？
• RQ2 与 OVE 等现有方法相比，所提出的 a&r 方法是否能提供更紧的边际似然下界？
• RQ3 在大规模分类任务中，a&r 是否能在对数似然和准确率方面优于精确 softmax 和 OVE？
• RQ4 在包含最多 13,000 个类别的数据集上，a&r 在多项式 probit 和多项式逻辑模型中的表现如何？
• RQ5 在每轮训练的运行时间方面，a&r 相较于 OVE 和精确 softmax 的计算效率如何？

主要发现

• 在除 Bibtex 外的所有数据集上，softmax a&r 的证据下界（ELBO）均优于 OVE，而在 Bibtex 上其下界也几乎与精确 softmax 相当。
• 在 MNIST 和 Bibtex 数据集上，a&r 的 ELBO 几乎与通过最大似然估计获得的精确边际似然无法区分。
• 在除 EURLex-4K 外的所有数据集上，softmax a&r 在测试对数似然和准确率方面均优于 OVE，而在 EURLex-4K 上 OVE 略优。
• 对于多项式 probit 和逻辑模型，a&r 在 Omniglot 和 Bibtex 上的预测性能优于 OVE，且多项式 probit a&r 在 EURLex-4K 和 AmazonCat-13K 上优于所有方法。
• 对于 softmax 模型，a&r 每轮训练的运行时间几乎与 OVE 相同；对于多项式模型，仅略高，表明其具有出色的计算效率。
• 即使在类别数高达 13,000（如 AmazonCat-13K）时，该方法仍能保持高预测准确率，证实其在真实世界大规模词汇任务中的可扩展性。