[論文レビュー] 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 操作を用いることで、学習と推論の高速化を実現する新しいトレーサブル確率的回路（EiNets）を提案する。効率的なテンソル結合と自動微分を活用することで、従来手法に比べて最大100倍の高速化とメモリ削減を達成し、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 アーキテクチャを設計し、速度とメモリ効率の両方を向上させること。

提案手法

• EiNets は、複数の和および積演算を1つの高度に最適化されたテンソル結合に統合したモノリシックな einsum 操作の階層を用いて確率的回路を表現する。
• 訓練中に数値安定性を保つために、標準的な log-sum-exp ティックを einsum の文脈に拡張した「log-einsum-exp」のテクニックを導入する。
• 複数の子を持つ和ノードについては、2段階のプロセスに分解する：まず einsum を用いて個別の和演算を計算し、その後、専用の「ミキシングレイヤー」を用いて要素ごとの凸混合を適用する。
• 多子の和ノードを段階的な単純な和ノードにオーバーパrameter化し、その後にミキシングレイヤーを配置することで、非二分木を含む任意の PC 構造をサポートする。
• 期待値最大化（EM）の実装を簡素化するために、自動微分を活用し、対数尤度の勾配を直接計算することで、学習を高速化する。
• 推論と学習を自然領域で実行することで、従来の PC フレームワークが採用する対数領域計算の性能劣化を回避する。

実験結果

リサーチクエスチョン

• RQ1 モノリシックな einsum 操作は、確率的回路の推論および学習における速度とメモリ効率を顕著に向上させることができるか？
• RQ2 EiNets は、従来のトレーサブル確率的モデルでは到達不可能だった、SVHN や CelebA といった大規模な画像データセットにスケーラブルに適用可能か？
• RQ3 自動微分の使用は、従来のアプローチと比較して、確率的回路における EM の実装を簡素化するか？
• RQ4 LibSPN や SPFlow といった既存の PC フレームワークと比較して、EiNets のトレーニング時間、推論速度、メモリ消費量における性能はどのように異なるか？

主な発見

• EiNets は、LibSPN や SPFlow といった従来の実装と比較して、最大2桁の速度向上とメモリ削減を達成しており、特に大規模データセットにおいて顕著である。
• モデルは SVHN や CelebA で正常に学習が可能であり、これらの複雑なデータセットに対しても、EiNets が忠実な生成画像モデルとして機能できることを示している。
• EiNets の推論時間は SPFlow より最大3桁の高速化を達成しており、ネットワークの深さが高くなっても顕著な改善が見られる。
• レイヤーごとに1つの einsum 操作を用いることで、CPU および GPU における計算が極めて効率的になり、ミキシングレイヤーにおけるテンソルのリシェイピングやゼロパディングに起因するオーバーヘッドが最小限に抑えられる。
• log-einsum-exp テクニックにより、訓練中に数値安定性が保たれ、性能劣化を伴わずに自然領域での最適化が信頼性を持って可能になる。
• 自動微分により、手動での勾配導出の必要性が排除され、学習プロセスがより堅牢で実装が容易になる。