HyperSPNs: Compact and Expressive Probabilistic Circuits

1 Introduction . One of the core motivations for building models of probability distributions from data is to reason about the data distribution . For example , given a model of the distribution over driving routes , we may like to reason about the probability that a driver will take a certain route given that there is congestion on a certain street . Or , given a distribution over weather conditions , we may like to reason about the probability of it raining k times in the next year . Traditional probabilistic models such as Categorical/Gaussian Mixture Models or Hidden Markov Models ( HMMs ) are wellequipped to answer these types of queries with exact probabilities that are consistent with the model distribution [ 1 ] . However , their learned model distribution might not accurately approximate the data distribution due to their simplicity . On the other end of the spectrum , the modern wave of deep generative models has largely focused on learning accurate approximations of the true data distribution , at the cost of tractability . High capacity models such as autoregressive or flow models can mimic the true data distribution with great fidelity [ 33 , 13 , 22 ] . However , they are not designed for reasoning about the probability distribution beyond computation of likelihoods . More illustrative of this trend are GANs and EBMs , which are expressive enough to sample high-quality images , but give up even the ability to compute exact likelihoods [ 9 , 14 ] . To maintain the ability to reason about the distribution induced by our model , we need to explore within tractable probabilistic models . Of these model families , probabilistic circuits ( e.g . Sum Product Networks ) are one of the most expressive and general , subsuming shallow mixture models , and HMMs while maintaining their tractable properties [ 6 , 26 , 34 ] . Their expressive power comes from their hierarchical / deep architecture , allowing them to express a large number of modes in their distribution . Their tractability comes from global constraints imposed in their network structure , enabling efficient and exact computation of likelihoods , marginals , and more . Probabilistic circuits 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ( PC ) are a popular choice for density estimation [ 35 , 18 ] and approximate inference in discrete settings [ 30 ] . Given the status of probabilistic circuits as one of the most expressive models among tractable probabilistic model families , recent works have looked into pushing the limits of expressivity of probabilistic circuits [ 15 , 25 , 24 ] . Naturally , the bigger and deeper the probabilistic circuit , the greater the expressiveness of the model . This has led to a trend of building large probabilistic circuits via design choices such as tensorizing model parameters , forgoing structure learning , and more . For example , [ 24 ] report the training of probabilistic circuits with 9.4M parameters . However , larger probabilistic circuits are more susceptible to overfitting . In addition , the so-called double-descent phenomenon [ 19 ] ( wherein highly overparameterized models exhibit low generalization gap ) has yet to be observed in probabilistic circuits . This puts even more importance on effective regularization strategies for training large probabilistic circuits . Recent works have used dropout [ 25 ] and weight decay [ 37 ] , but there has not been much exploration of other regularization strategies . In this work , we propose HyperSPNs , which regularize by aggressively limiting the degrees of freedom of the model . Drawing inspiration from HyperNetworks [ 11 ] , we generate the weights of the probabilistic circuit via a small-scale external neural network . More precisely , we partition the parameters ( mixture weights ) of the PC into several sectors . For each sector , we learn a lowdimensional embedding , then map that to the parameters of the sector using the neural network . The generated parameters still structurally form a PC , so we retain the same ability to reason about the probability distribution induced by the HyperSPN . HyperSPNs combine the greater expressiveness of large probabilistic circuits with the better generalization of models with fewer degrees of freedom . The external neural network has much fewer parameters than the original PC , effectively regularizing the PC through a soft weight-sharing strategy . As a bonus , the memory requirement for storing the model is much smaller , and the memory requirement for evaluating the model can also be drastically reduced by materializing sectors of the PC “ on the fly ” . We verify the empirical performance of HyperSPNs on density estimation tasks for the Twenty Datasets benchmark [ 12 ] , the Amazon Baby Registries benchmark [ 8 ] , and the Street View House Numbers ( SVHN ) [ 20 ] dataset , showing generalization improvements over competing regularization strategies in all three settings . Our results suggest that HyperSPNs are a promising framework for training probabilistic circuits . 2 Background . Probabilistic circuits ( PCs ) are a family of generative models known for their tractability properties [ 6 , 26 , 34 ] . They support efficient and exact inference routines for computing likelihoods and marginals/conditionals of their probability distribution . The internals of a PC consist of a network ( edge-weighted Directed Acyclic Graph ) of sum and product nodes , stacked in a deep/hierarchical manner and evaluated from leaves to root . PCs gain their tractability by enforcing various structural constraints on their network , such as decomposability , which restricts children of the same product node to have disjoint scopes [ 5 , 6 ] . Different combination of constraints leading to different types of PCs , including Arithmetic Circuits [ 6 ] , and most notably Sum Product Networks ( SPNs ) [ 26 ] . The probability of an input x on the PC is defined to be the value groot ( x ) obtained at the root of the PC after evaluating the network topologically bottom-up . Internally , edges leading into a sum node are weighted with mixture probabilities α ( that sum to 1 ) , and edges leading into a product node are unweighted . Leaf nodes at the base of the network contain leaf distributions l. The nodes of the PC are evaluated as follows , with ch ( i ) denoting the children of node i : gi ( x ) = li ( x ) i is leaf node∑ j∈ch ( i ) αijgj ( x ) i is sum node∏ j∈ch ( i ) gj ( x ) i is product node As long as the leaf distributions are tractable and the network structure satisfies the constraints of decomposability and smoothness , the output at groot is guaranteed to be normalized , and corresponding marginals/conditionals can be evaluated with linear-time exact inference algorithms [ 5 , 6 ] . The network structure of a PC refers to the set of nodes and edge connections in the DAG , excluding the edge weights . There are various ways of constructing the network structure of PCs , e.g. , through X1 X1 X2 X2 × × × X3 X3 X4 X4 × × × × × × ( a ) We show the structure of a RAT-SPN over n = 4 variables using layer size k = 3 . Solid edges are weighted with trainable parameters , and dashed edges are unweighted . Each fully connected layer of solid edges can be treated as a sector . ( b ) Counting the # of parameters , given the number of variables n and layer size k. Figure 1 : Given a RAT-SPN , we can partition its set of parameters into sectors . The trainable parameters ( solid edges ) in Figure 1a can be grouped into fully-connected layers of size 1× k , k × k , and k×2 , oriented in a binary-tree structure . In Figure 1b , we visualize the sectors as blocks , counting the number of blocks in each layer of the binary tree to arrive at the total parameter count . The sector abstraction is a key ingredient in implementing the HyperSPN ’ s soft weight-sharing strategy . the use of domain knowledge or through search-based methods [ 6 , 17 , 7 , 28 ] . Recently , there has been a trend of abstracting away individual sum and product nodes , and instead dealing directly with layers that satisfy the structural constraints necessary for tractability [ 25 , 24 , 30 ] . Prescribing a network structure through a stacking of these layers leads to memory and hardware efficiencies that allow us to scale up the training of PCs to millions of edges . Given a PC with a fixed network structure , we can learn its edge weights to best fit the probability distribution of data . We will refer to the full set of mixture probabilities leading into sum nodes in the PC as its parameters . These parameters can be learned using methods such as ExpectationMaximization ( EM ) [ 23 ] , signomial program formulations [ 38 ] , or simply via ( projected ) gradient descent . EM has been a well-studied choice for optimizing PCs [ 23 , 24 ] , although it requires handderived updates that can be tedious and sometimes even incorrect [ 26 ] . In this work , we rely on gradient descent , which is straightforward and works out-of-the-box on new families of models [ 37 ] . Learning a PC with a large number of parameters can be prone to overfitting . One regularization strategy that has been proposed is “ probabilistic dropout ” [ 25 ] , which takes inspiration from the standard dropout technique in neural networks [ 31 ] and randomly ignores a subset of edges leading into sum nodes . Weight decay has also been commonly used as another form of regularization [ 37 ] . Unfortunately , there has been little exploration done around regularization strategies , and some current approaches only work in limited settings ( e.g. , dropout has been limited to settings that use PCs as a discriminative classifier [ 25 ] ) . 2.1 RAT-SPNs . Two particular PC structures have been recently proposed , namely RAT-SPNs [ 25 ] and EiNETs [ 24 ] , both of which scale to millions of parameters . They are conceptually similar , the main difference being that EiNETs are designed for better GPU performance . To ground our discussion around regularizing large PCs , we will describe our methods just on RAT-SPNs , although we also apply our methods on EiNETs in later experiments too . In RAT-SPNs , the input variables are organized into a binary tree , with one leaf in the binary tree corresponding to each variable . We start from the base of the SPN and construct its full structure by stacking sum layers and product layers . For each leaf node in the binary tree , we create one fully connected sum layer on top of the SPN leaves . For each internal node in the binary tree , we create one product layer that merges its two children layers in the SPN , and then stack one fully connected sum layer on top . We illustrate this construction in Figure 1 . Finally , RAT-SPNs create so-called replicas of this structure , by repeating the construction but with permuted orderings of input variables in the binary tree leaves . The replicas are then combined under one top-level mixture distribution ( for readability , Figure 1 does not show any replicas ) . The size of each sum layer in the RAT-SPN is determined by a single constant k. For convenience , we also let the size of each product layer be k. This fully specifies the structure of our SPN . For the rest of the paper , it is helpful to abstractly group the parameters of the RAT-SPN into disjoint sectors . Corresponding to the root of the binary tree , we have a sector with 1×k parameters . Corresponding to each internal node in the binary tree , we have a sector with k × k parameters . Finally , corresponding to each leaf in the binary tree , we have a sector with k × l parameters , where l is the number of SPN leaf distributions per variable ( we use l = 2 unless otherwise stated ) . As shown in Figure 1 , for n variables ( where n is a power of 2 ) and r replicas , this corresponds to 2n − 1 sectors totaling O ( rnk2 ) parameters . In particular , for each replica there are ( n − 2 ) k2 + ( 2n+ 1 ) k parameters . For example , in a problem domain with 1000 variable , using layer size k = 10 and replica r = 10 leads to a RAT-SPN of around 1M parameters . 3 HyperSPNs We propose HyperSPNs , an approach of regularizing probabilistic circuits with a large number of parameters by limiting the degrees of freedom of the model using an external neural network . We describe our method concretely on RAT-SPNs , which we will refer to as simply SPNs . To regularize the SPN , we introduce a smaller-scale external neural network to generate the parameters of the SPN ( see Figure 2 ) . We associate each of the sectors of the SPN with a sector embedding . For each sector , we map the embedding to the parameters of the sector via the neural network . Then , we use these mapped parameters directly as the parameters of the SPN , re-normalize such that parameters leading into a sum node add up to 1 , and evaluate the SPN as usual . We visualize the computation graph of this process in Figure 3 . We materialize the SPN in the forward pass , and propagate the gradients through the SPN in the backward pass . We optimize the neural network and embeddings jointly using gradient descent . HyperSPNs can aggressively reduce the degrees of freedom of the original SPN . Using embeddings of size h , the number of learnable parameters is reduced to O ( ( rn + k2 ) h ) . For the same above example with n = 1000 , k = 10 , r = 10 , choosing embeddings of dimension h = 5 reduces the number of learnable parameters 20 fold , from around 1M down to around 50k parameters . The lower degrees of freedom mitigates overfitting of the model . At the same time , the materialized SPN has a large number of mixture weights , and thus can express distribution families with a greater number of clusters compared to a small SPN with the same degrees of freedom . We visualize this with a motivating example in Figure 4 , where we represent a dataset using a grid that is split recursively for each variable ( we only draw the splits for the first 2 variables ; the rest is implied ) . Intuitively , the number of parameters of the SPN is tied closely to the number of clusters in the distribution . As a result , training a large SPN can produce many clusters that overfit to the training points in blue ( Figure 4a ) , and training a small SPN can lead to coarse factorizations and inflexible clusters ( Figure 4b ) . Both cases can lead to poor generalization on the testing points in orange . On the other hand , the HyperSPN uses a soft weight-sharing strategy by storing its parameters as low-dimensional embeddings and decoding them through a small-scale shared neural network . This allows the model to learn a high number of clusters without suffering from too high of a degree of freedom . In Figure 4c , we visualize the soft weight-sharing as fixing the shape of the cluster-triples within each quadrant , while still giving the model control over the center position of each cluster-triple . In this way , HyperSPNs can improve generalization compared to a large SPN with the same number of mixture weights and a small SPN with the same degrees of freedom . We verify this phenomenon on a hand-crafted dataset in the Appendix , and further observe similar improvements on real world datasets in our experiments .

Summary: The paper proposes a novel method for regularizing large scale probabilistic circuits: The PC's parameters are partitioned into evenly sized sectors, and each sector is represented via a low dimensional embedding. A shared neural network is used to map the embeddings in to the PC's original parameter space. Experimental results show that this reduction in the number of trainable parameters yields improved generalization performance when compared to previous methods, mainly weight decay.