Source-linked AI summary
Einsum Networks: Fast and Scalable Learning of Tractable Probabilistic Circuits
Robert Peharz, Steven Lang, Antonio Vergari, Karl Stelzner, Alejandro Molina, Martin Trapp, Guy Van den Broeck, Kristian Kersting, Zoubin Ghahramani
TL;DR
Probabilistic circuits support exact and efficient inference, but deep-learning-style implementations remain difficult to train and scale on real-world data. The paper introduces Einsum Networks, which consolidate PC arithmetic into large einsum operations and simplify EM training, achieving substantial efficiency gains and scaling to SVHN and CelebA. EiNets also support high-quality generative image modeling with tractable conditional inference.
Problem
Deep-learning-style probabilistic-circuit implementations remain difficult to train on real-world data because their computational graphs are sparsely connected.
Method
Einsum Networks combine product and sum operations into monolithic einsum layers and implement EM through automatic differentiation with stochastic online EM.
Results
EiNets train generative models on SVHN and CelebA, provide high-quality image samples with tractable conditional inference, and improve training and evaluation speed and memory use by up to one or two orders of magnitude.
Takeaways & Limitations
EiNets make probabilistic circuits more scalable while retaining tractable inference for generative image tasks such as inpainting.
Takeaways & Limitations
Stochastic EM introduces step-size and batch-size hyperparameters and does not guarantee increasing training likelihood at every iteration.
Abstract
from arXiv · showhide
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.
1. Introduction
Probabilistic circuits target exact, efficient inference, addressing difficulties that deep neural probabilistic models face on complex conditional queries. EiNets improve PC scalability by consolidating layered arithmetic into einsum operations and by simplifying training with automatic differentiation and stochastic EM.
- Motivation: Deep neural probabilistic models offer strong representation and learning on large datasets, but their inference is often approximate, biased, unreliable, or computationally expensive.Monte Carlo inference becomes unreliable and expensive beyond simple queries, while variational inference can be biased and difficult to analyze.
- Motivation: Conditional densities must predict query variables from evidence while marginalizing missing variables, yet evaluating such queries is notoriously hard for common generative models.The query, missing, and evidence variable sets may be chosen arbitrarily.
- Probabilistic circuits: Probabilistic circuits represent densities with mixtures, factorizations, and tractable leaves, while decomposability enables integrals to be computed in linear time in circuit size.These structural properties support exact and efficient inference.
- Einsum Networks: EiNets compute product and sum operations on each topological layer with a single monolithic einsum operation.This lifts the main computation into parallel CPU- and GPU-efficient operations; the log-einsum-exp trick supports numerical stability.
- Einsum Networks: EiNets use a hierarchy of large einsum layers and therefore are called Einsum Networks.The design targets faster evaluation and training of probabilistic circuits.
- Training: The paper simplifies PC training by implementing EM through gradients of model log-probability and uses stochastic online EM for faster learning or large datasets.Experiments train generative models on SVHN and CelebA, datasets previously difficult for PCs.
- Applications: EiNets produce high-quality image samples while retaining tractable inference for conditional densities, supporting image inpainting and other inference tasks.The paper presents these capabilities as applications of the model's tractable generative structure.
2. Probabilistic Circuits
Probabilistic circuits are structured computational graphs that combine tractable leaves with sum and product nodes. Smoothness and decomposability provide the probabilistic interpretation and efficient marginalization properties emphasized for the circuits studied here.
- PCs: Probabilistic circuits include arithmetic circuits, sum-product networks, cutset networks, and probabilistic sentential decision diagrams.They form a family of models designed to support exact and efficient inference routines.
- Definition: A probabilistic circuit is a directed acyclic graph with a scope function, probability-density leaves, and internal sum or product nodes.The scope function assigns each node a subset of variables and equals the union of child scopes for internal nodes.
- Node operations: Sum nodes compute convex combinations of their children, while product nodes compute products of their children.Sum weights are nonnegative and normalized.
- Interpretation: PCs resemble neural networks whose leaves compute densities and whose internal nodes compute weighted sums or multiplicative interactions over scoped variables.The output is typically taken from a root node with full scope.
- Tractability: Without structural constraints, a PC defines a density but does not yet permit tractable inference because its normalization constant can be hard to compute.Further constraints are required to unlock tractable inference routines.
- Decomposability: Decomposability requires product-node children to have pairwise non-overlapping scopes, making relevant nested integrals computable in linear time in circuit size.This includes the normalization constant.
- Smoothness: Smoothness requires sum-node children to have identical scopes and gives nodes a well-defined probabilistic interpretation in smooth decomposable PCs.Under these conditions, nodes are properly normalized densities.
- Scope: The paper studies smooth and decomposable PCs, while excluding determinism and structured decomposability from consideration.Determinism enables exact probability maximization, whereas structured decomposability is a stronger structural property.
3. Einsum Networks
EiNets vectorize probabilistic-circuit nodes and evaluate product-sum computations with numerically stable einsum operations. They organize entire layers for parallel computation and support efficient EM-oriented leaf parameterization.
- 3.1. Vectorizing Probabilistic Circuits: Vectorized leaves contain K densities, while product nodes compute outer products over all child combinations.The paper uses a common vector length K for all leaves and sums.
- 3.2. The Basic Einsum Operation: The basic einsum operation computes a vectorized product followed by weighted summation using a non-negative K × K × K weight tensor.The tensor indices label the two child vectors and the output sum vector; repeated indices multiply and omitted indices are summed.
- 3.2. The Basic Einsum Operation: Previous implementations explicitly construct products and use broadcasted log-domain operations, whereas EiNets express the same multilinear computation directly with einsum.The prior workaround introduces overhead and requires explicit product allocation.
- 3.2. The Basic Einsum Operation: EiNets retain probabilistic values in the log-domain while keeping weights in the linear domain to avoid numerical underflow.A log-einsum-exp-style computation is stable when sum weights are positive; positivity can be enforced by smoothing or lower bounds.
- 3.3. The Einsum Layer: EiNets parallelize complete sum-product layers by adding a layer index to the einsum operation and use two consecutive sum layers for multi-child sums.The second layer is a mixing layer, representing an over-parameterization of the original multi-child sums.
- 3.4. Input Nodes: EiNet leaves compute exponential-family log-densities, with replica indexing ensuring each leaf has K private distributions over disjoint scopes.Expectation-form parameters are retained to facilitate EM learning.
4. Experiments
Experiments compare EiNets with LibSPN and SPFlow on synthetic probabilistic circuits and image datasets. EiNets are substantially faster and more memory-efficient, while producing plausible SVHN and CelebA samples and conditional reconstructions.
- 4.1. Efficiency Comparison: EiNets were compared with LibSPN and SPFlow on randomized binary PC trees, varying depth D, replica count R, and vector length K.The experiments measured training time, memory consumption, and inference time using synthetic Gaussian data.
- 4.1. Efficiency Comparison: One to two orders of magnitude faster training was observed for EiNets than for competitors, especially on large models.The comparison used a log-scale plot of training time and GPU-memory consumption.
- 4.1. Efficiency Comparison: An order of magnitude lower memory consumption was observed for EiNets at large K than for LibSPN or SPFlow.The authors attribute this to avoiding explicit storage of product nodes.
- 4.1. Efficiency Comparison: One to two orders of magnitude faster inference was observed for EiNets on large models than for the other implementations.Inference results were deferred to the supplementary material.
- 4.2. EiNets as Generative Image Models: EiNets were trained as generative models on SVHN 32 × 32 RGB images and CelebA 128 × 128 RGB face images.The image experiments used image-tailored PD structures and mixtures of 100 cluster-specific EiNets.
- 4.2. EiNets as Generative Image Models: SVHN samples were described as compelling, CelebA samples as somewhat over-smoothed but capturing facial structure, and both datasets showed stripy PD-architecture artifacts.The samples were generated using the means of Gaussian leaves.
- 4.2. EiNets as Generative Image Models: Conditional sampling produced plausible image inpaintings from visible image halves, particularly for SVHN.The reconstructions were conditioned on visible portions of test images.
5. Conclusion
The paper addresses probabilistic circuits’ scalability by reducing training time and memory use by one or two orders of magnitude. It presents these improvements as a basis for further development of tractable models.
- 5. Conclusion: EiNets address probabilistic circuits’ scalability obstacle with training-speed and memory-use improvements of one or two orders of magnitude.The conclusion frames these improvements relative to the scalability of PCs compared with unconstrained models.
A. Organizing EiNets in Topological Layers
Algorithm 1 organizes an EiNet graph into topologically ordered layers containing leaves, sums, or products. It repeatedly adds eligible sums and products whose parents have already been visited.
- A. Organizing EiNets in Topological Layers: The layer-organization algorithm takes an EiNet computational graph and constructs a topologically ordered list of pure node layers.Each layer contains exclusively leaves, sums, or products.
- A. Organizing EiNets in Topological Layers: The algorithm repeatedly selects unvisited sum nodes whose parents are visited, then performs the analogous operation for product nodes.The root sum is included in the first iteration when applicable.
- A. Organizing EiNets in Topological Layers: All leaf nodes are inserted as the bottom layer, and every node receives inputs only from layers with smaller indices.This establishes the returned list as topologically sorted.
B. The Mixing Layer
The mixing layer converts sum nodes with multiple children into two consecutive vectorized sum layers. It preserves the represented linear functions while introducing tensor-construction and zero-padding overhead.
- B. The Mixing Layer: The decomposition introduces one simple sum node per child, followed by aggregated sums that combine the simple-sum results.This construction is illustrated for sum nodes with three and two children.
- B. The Mixing Layer: A multi-child sum layer is decomposed into a standard einsum layer followed by a mixing layer that performs element-wise mixtures.The first layer uses simple sums, while the second aggregates their outputs.
- B. The Mixing Layer: The construction preserves the original sum nodes’ linear functions but incurs copy overhead and wasted computation from zero padding.The overhead is exchanged for support of a wider range of PC structures than random binary trees.
- B. The Mixing Layer: The mixing layer stores K-dimensional probability vectors in a D × M × K tensor and computes a convex combination over the child dimension.The child dimension is zero-padded for sums with fewer than D children.
- B. The Mixing Layer: Simple sum layers skip the mixing-layer construction when the original sums have only one child.The extra layer is needed only for sums with multiple children.
C. Inference Time Comparison
EiNets achieve substantially faster inference than LibSPN and SPFlow, with speedups reaching three orders of magnitude at maximal depth.
- Up to three orders of magnitude faster inference is observed for EiNets versus SPFlow at maximal depth.Inference time is measured per sample using batches of 100 test samples.
D. On Using the Mean When Generating Images (Addendum, October 2025)
For image generation, the procedure turns off Gaussian leaf noise while retaining hierarchical sampling of probabilistic-circuit sum nodes. This is justified for visual generation because pixelwise Gaussian noise degrades image quality, although it is not appropriate when probabilistic reasoning is the goal.
- Image-generation procedure: Gaussian leaf noise is turned off during image generation, while sum nodes are sampled hierarchically to select leaves covering all modeled pixels.Exact sampling from P(x) would additionally sample each leaf distribution and concatenate the resulting pixel values.
- Relation to VAEs: The procedure parallels turning off decoder noise in variational autoencoders during visualization.Both models are hierarchical latent-variable models, with discrete sum-node selections in PCs corresponding to continuous latent variables in VAEs.
- Motivation and scope: Under pixelwise independence in p(x | z), removing Gaussian leaf noise avoids independent pixel noise that would degrade visual quality.This practice is justified for generative modeling, but not when probabilistic reasoning requires sampling from the full distribution.