Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic\n Circuits

Abstract

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