[论文解读] Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuits

Robert Peharz, Steven Lang+7 more|arXiv (Cornell University)|Apr 13, 2020

Machine Learning and Data Classification参考文献 34被引用 27

一句话总结

该论文提出了爱因斯坦网络（Einsum Networks, EiNets），这是一种新型可 tractable 概率电路的实现方式，通过用单一的 einsum 操作替代稀疏的计算图，显著加速了学习与推理过程。通过利用高效的张量收缩和自动微分，EiNets 在先前方法的基础上实现了高达两个数量级的速度提升与内存减少，使大规模数据集（如 SVHN 和 CelebA）上的可扩展训练成为可能，同时支持对复杂条件查询的精确高效推理。

ABSTRACT

Probabilistic circuits (PCs) are a promising avenue for probabilistic modeling, as they permit a wide range of exact and efficient inference routines. Recent ``deep-learning-style'' implementations of PCs strive for a better scalability, but are still difficult to train on real-world data, due to their sparsely connected computational graphs. In this paper, we propose Einsum Networks (EiNets), a novel implementation design for PCs, improving prior art in several regards. At their core, EiNets combine a large number of arithmetic operations in a single monolithic einsum-operation, leading to speedups and memory savings of up to two orders of magnitude, in comparison to previous implementations. As an algorithmic contribution, we show that the implementation of Expectation-Maximization (EM) can be simplified for PCs, by leveraging automatic differentiation. Furthermore, we demonstrate that EiNets scale well to datasets which were previously out of reach, such as SVHN and CelebA, and that they can be used as faithful generative image models.

研究动机与目标

• 解决现有概率电路（PC）实现因稀疏连接的对数域计算图而导致的可扩展性差和计算成本高的问题。
• 实现对大规模真实世界数据集（如 SVHN 和 CelebA）的高效训练与精确推理，这些数据集此前对 PC 模型而言是不可行的。
• 通过利用自动微分，简化基于 EM 的 PC 学习，降低实现复杂度。
• 设计一种新的 PC 架构，通过高效的张量操作统一乘积与求和操作，从而同时提升速度与内存效率。

提出的方法

• EiNets 使用分层的单一 einsum 操作表示概率电路，将多个求和与乘积操作合并为单一高度优化的张量收缩操作。
• 模型采用一种新颖的 'log-einsum-exp' 技巧，在训练过程中保持数值稳定性，将标准的 log-sum-exp 技巧扩展至 einsum 设置。
• 对于具有多个子节点的求和节点，EiNets 将其分解为两阶段过程：首先通过 einsum 计算各个求和操作，然后通过专用的 '混合层' 应用逐元素凸组合。
• 通过将多子节点求和节点过度参数化为级联的简单求和节点，再接一个混合层，该架构支持任意 PC 结构，包括非二叉树结构。
• 通过自动微分加速训练，直接计算对数似然的梯度，从而简化了期望最大化（EM）的实现。
• 推理与学习在自然域中进行，避免了先前 PC 框架中使用的对数域计算带来的性能损失。

实验结果

研究问题

• RQ1 单一 einsum 操作是否能显著提升概率电路推理与训练的速度和内存效率？
• RQ2 EiNets 是否能够扩展到 SVHN 和 CelebA 等大规模图像数据集，这些数据集此前对可 tractable 概率模型而言难以企及？
• RQ3 与传统方法相比，使用自动微分是否能简化概率电路的 EM 实现？
• RQ4 在训练时间、推理速度和内存消耗方面，EiNets 与现有 PC 框架（如 LibSPN 和 SPFlow）相比表现如何？

主要发现

• 与先前实现（如 LibSPN 和 SPFlow）相比，EiNets 在大规模数据集上实现了高达两个数量级的速度提升与内存减少。
• 该模型成功在 SVHN 和 CelebA 上完成训练，证明了 EiNets 尽管面对复杂数据集，仍可作为忠实的生成式图像模型。
• EiNets 的推理时间比 SPFlow 快 up to 三个数量级，即使在高网络深度下也表现出显著提升。
• 每层仅使用一个 einsum 操作，使得在 CPU 和 GPU 上都能实现高度高效的计算，混合层中的张量重塑与零填充开销极低。
• log-einsum-exp 技巧确保了训练过程中的数值稳定性，使自然域中的可靠优化成为可能，且无性能损失。
• 自动微分通过消除手动梯度推导的需要，简化了 EM 训练，使学习过程更加稳健且易于实现。